From 12d451349e0d8178447d23d6cfce24d04cb67b52 Mon Sep 17 00:00:00 2001
From: Jon Marius Venstad <jvenstad@yahoo-inc.com>
Date: Fri, 14 Jun 2019 15:34:38 +0200
Subject: Allow Endpoint to send arbitrary, authenticated requests

---
 .../src/main/java/ai/vespa/hosted/cd/Endpoint.java |  4 ++++
 .../java/ai/vespa/hosted/cd/http/HttpEndpoint.java | 22 ++++++++++++++++------
 2 files changed, 20 insertions(+), 6 deletions(-)

(limited to 'tenant-cd')

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;
@@ -57,16 +59,24 @@ public class HttpEndpoint implements TestEndpoint {
         return endpoint;
     }
 
+    @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()));
-- 
cgit v1.2.3