Only wait for requested service if there are multiple

2 files changed, 36 insertions, 19 deletions

diff --git a/client/go/internal/cli/cmd/status_test.go b/client/go/internal/cli/cmd/status_test.go
index 4ed13a7ef22..7e1e266f694 100644
--- a/client/go/internal/cli/cmd/status_test.go
+++ b/client/go/internal/cli/cmd/status_test.go
@@ -33,6 +33,7 @@ func TestStatusCommandMultiCluster(t *testing.T) {
 	client := &mock.HTTPClient{}
 	cli, stdout, stderr := newTestCLI(t)
 	cli.httpClient = client
+	cli.retryInterval = 0
 
 	mockServiceStatus(client)
 	assert.NotNil(t, cli.Run("status"))
@@ -50,6 +51,18 @@ Container foo at http://127.0.0.1:8080 is ready
 	assert.Equal(t, "Container foo at http://127.0.0.1:8080 is ready\n", stdout.String())
 }
 
+func TestStatusCommandMultiClusterWait(t *testing.T) {
+	client := &mock.HTTPClient{}
+	cli, _, stderr := newTestCLI(t)
+	cli.httpClient = client
+	cli.retryInterval = 0
+	mockServiceStatus(client, "foo", "bar")
+	client.NextStatus(400)
+	assert.NotNil(t, cli.Run("status", "--cluster", "foo", "--wait", "10"))
+	assert.Equal(t, "Waiting up to 10s for cluster discovery...\nWaiting up to 10s for container foo...\n"+
+		"Error: unhealthy container foo after waiting up to 10s: status 400 at http://127.0.0.1:8080/ApplicationStatus: aborting wait: got status 400\n", stderr.String())
+}
+
 func TestStatusCommandWithUrlTarget(t *testing.T) {
 	assertStatus("http://mycontainertarget:8080", []string{"-t", "http://mycontainertarget:8080"}, t)
 }
diff --git a/client/go/internal/cli/cmd/waiter.go b/client/go/internal/cli/cmd/waiter.go
index a1919de798d..1dc227cd683 100644
--- a/client/go/internal/cli/cmd/waiter.go
+++ b/client/go/internal/cli/cmd/waiter.go
@@ -47,7 +47,7 @@ func (w *Waiter) Service(target vespa.Target, cluster string) (*vespa.Service, e
 	if targetType.url != "" && cluster != "" {
 		return nil, fmt.Errorf("cluster cannot be specified when target is an URL")
 	}
-	services, err := w.Services(target)
+	services, err := w.services(target)
 	if err != nil {
 		return nil, err
 	}
@@ -55,39 +55,43 @@ func (w *Waiter) Service(target vespa.Target, cluster string) (*vespa.Service, e
 	if err != nil {
 		return nil, errHint(err, "The --cluster option specifies the service to use")
 	}
-	if w.Timeout > 0 {
-		w.cli.printInfo("Waiting up to ", color.CyanString(w.Timeout.String()), " for ", color.CyanString(service.Description()), " to become available...")
-	}
-	if w.wait() {
-		if err := service.Wait(w.Timeout); err != nil {
-			return nil, err
-		}
+	if err := w.maybeWaitFor(service); err != nil {
+		return nil, err
 	}
 	return service, nil
 }
 
 // Services returns all container services available on target.
 func (w *Waiter) Services(target vespa.Target) ([]*vespa.Service, error) {
-	if w.Timeout > 0 {
-		w.cli.printInfo("Waiting up to ", color.CyanString(w.Timeout.String()), " for cluster discovery...")
-	}
-	services, err := target.ContainerServices(w.Timeout)
+	services, err := w.services(target)
 	if err != nil {
 		return nil, err
 	}
 	for _, s := range services {
-		if w.Timeout > 0 {
-			w.cli.printInfo("Waiting up to ", color.CyanString(w.Timeout.String()), " for ", s.Description(), "...")
-		}
-		if w.wait() {
-			if err := s.Wait(w.Timeout); err != nil {
-				return nil, err
-			}
+		if err := w.maybeWaitFor(s); err != nil {
+			return nil, err
 		}
 	}
 	return services, nil
 }
 
+func (w *Waiter) maybeWaitFor(service *vespa.Service) error {
+	if w.Timeout > 0 {
+		w.cli.printInfo("Waiting up to ", color.CyanString(w.Timeout.String()), " for ", service.Description(), "...")
+	}
+	if w.wait() {
+		return service.Wait(w.Timeout)
+	}
+	return nil
+}
+
+func (w *Waiter) services(target vespa.Target) ([]*vespa.Service, error) {
+	if w.Timeout > 0 {
+		w.cli.printInfo("Waiting up to ", color.CyanString(w.Timeout.String()), " for cluster discovery...")
+	}
+	return target.ContainerServices(w.Timeout)
+}
+
 // Deployment waits for a deployment to become ready, returning the ID of the converged deployment.
 func (w *Waiter) Deployment(target vespa.Target, id int64) (int64, error) {
 	if w.Timeout > 0 {