Allow Endpoint to send arbitrary, authenticated requests

2 files changed, 20 insertions, 6 deletions

diff --git a/tenant-cd/src/main/java/ai/vespa/hosted/cd/Endpoint.java b/tenant-cd/src/main/java/ai/vespa/hosted/cd/Endpoint.java
index efeb4214ebd..dbbb969efe2 100644
--- a/tenant-cd/src/main/java/ai/vespa/hosted/cd/Endpoint.java
+++ b/tenant-cd/src/main/java/ai/vespa/hosted/cd/Endpoint.java
@@ -3,6 +3,8 @@ package ai.vespa.hosted.cd;
 import ai.vespa.hosted.cd.metric.Metrics;
 
 import java.net.URI;
+import java.net.http.HttpRequest;
+import java.net.http.HttpResponse;
 
 /**
  * An endpoint in a Vespa application {@link Deployment}, which allows document and metrics retrieval.
@@ -16,6 +18,8 @@ public interface Endpoint {
 
     URI uri();
 
+    <T> HttpResponse<T> send(HttpRequest.Builder request, HttpResponse.BodyHandler<T> handler);
+
     Search search(Query query);
 
     Visit visit(Selection selection);
diff --git a/tenant-cd/src/main/java/ai/vespa/hosted/cd/http/HttpEndpoint.java b/tenant-cd/src/main/java/ai/vespa/hosted/cd/http/HttpEndpoint.java
index 02c34501dd2..17703d8fbab 100644
--- a/tenant-cd/src/main/java/ai/vespa/hosted/cd/http/HttpEndpoint.java
+++ b/tenant-cd/src/main/java/ai/vespa/hosted/cd/http/HttpEndpoint.java
@@ -13,6 +13,8 @@ import ai.vespa.hosted.cd.TestEndpoint;
 import ai.vespa.hosted.cd.Visit;
 import ai.vespa.hosted.cd.metric.Metrics;
 
+import java.io.IOException;
+import java.io.UncheckedIOException;
 import java.net.URI;
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
@@ -58,15 +60,23 @@ public class HttpEndpoint implements TestEndpoint {
     }
 
     @Override
+    public <T> HttpResponse<T> send(HttpRequest.Builder request, HttpResponse.BodyHandler<T> handler) {
+        try {
+            return client.send(authenticator.authenticated(request).build(), handler);
+        }
+        catch (IOException | InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
     public Search search(Query query) {
         try {
             URI target = endpoint.resolve(searchApiPath).resolve("?" + query.rawQuery());
-            HttpRequest request = authenticator.authenticated(HttpRequest.newBuilder()
-                                                                         .timeout(query.timeout().orElse(Duration.ofMillis(500))
-                                                                                       .plus(Duration.ofSeconds(1)))
-                                                                         .uri(target))
-                                               .build();
-            HttpResponse<byte[]> response = client.send(request, HttpResponse.BodyHandlers.ofByteArray());
+            HttpResponse<byte[]> response = send(HttpRequest.newBuilder(target)
+                                                            .timeout(query.timeout().orElse(Duration.ofMillis(500))
+                                                                          .plus(Duration.ofSeconds(1))),
+                                                 HttpResponse.BodyHandlers.ofByteArray());
             if (response.statusCode() / 100 != 2) // TODO consider allowing 504 if specified.
                 throw new RuntimeException("Non-OK status code " + response.statusCode() + " at " + target +
                                            ", with response \n" + new String(response.body()));