Support cluster discovery for all target types

1 files changed, 13 insertions, 5 deletions

diff --git a/client/go/internal/cli/cmd/test.go b/client/go/internal/cli/cmd/test.go
index abee760efbb..58c254ad6e8 100644
--- a/client/go/internal/cli/cmd/test.go
+++ b/client/go/internal/cli/cmd/test.go
@@ -79,7 +79,7 @@ func runTests(cli *CLI, rootPath string, dryRun bool, waitSecs int) (int, []stri
 		if err != nil {
 			return 0, nil, errHint(err, "See https://docs.vespa.ai/en/reference/testing")
 		}
-		context := testContext{testsPath: rootPath, dryRun: dryRun, cli: cli}
+		context := testContext{testsPath: rootPath, dryRun: dryRun, cli: cli, clusters: map[string]*vespa.Service{}}
 		previousFailed := false
 		for _, test := range tests {
 			if !test.IsDir() && filepath.Ext(test.Name()) == ".json" {
@@ -100,7 +100,7 @@ func runTests(cli *CLI, rootPath string, dryRun bool, waitSecs int) (int, []stri
 			}
 		}
 	} else if strings.HasSuffix(stat.Name(), ".json") {
-		failure, err := runTest(rootPath, testContext{testsPath: filepath.Dir(rootPath), dryRun: dryRun, cli: cli}, waitSecs)
+		failure, err := runTest(rootPath, testContext{testsPath: filepath.Dir(rootPath), dryRun: dryRun, cli: cli, clusters: map[string]*vespa.Service{}}, waitSecs)
 		if err != nil {
 			return 0, nil, err
 		}
@@ -216,9 +216,15 @@ func verify(step step, defaultCluster string, defaultParameters map[string]strin
 		if err != nil {
 			return "", "", err
 		}
-		service, err = target.Service(vespa.QueryService, time.Duration(waitSecs)*time.Second, 0, cluster)
-		if err != nil {
-			return "", "", err
+		ok := false
+		service, ok = context.clusters[cluster]
+		if !ok {
+			// Cache service so we don't have to discover it for every step
+			service, err = context.cli.service(target, cluster, time.Duration(waitSecs)*time.Second)
+			if err != nil {
+				return "", "", err
+			}
+			context.clusters[cluster] = service
 		}
 		requestUrl, err = url.ParseRequestURI(service.BaseURL + requestUri)
 		if err != nil {
@@ -474,6 +480,8 @@ type testContext struct {
 	lazyTarget vespa.Target
 	testsPath  string
 	dryRun     bool
+	// Cache of services by their cluster name
+	clusters map[string]*vespa.Service
 }
 
 func (t *testContext) target() (vespa.Target, error) {