Add support for dynamic headers

2 files changed, 63 insertions, 2 deletions

diff --git a/vespa-http-client/src/test/java/com/yahoo/vespa/http/client/config/ConnectionParamsTest.java b/vespa-http-client/src/test/java/com/yahoo/vespa/http/client/config/ConnectionParamsTest.java
index 49ffabbf1d0..cb9270fdd3f 100644
--- a/vespa-http-client/src/test/java/com/yahoo/vespa/http/client/config/ConnectionParamsTest.java
+++ b/vespa-http-client/src/test/java/com/yahoo/vespa/http/client/config/ConnectionParamsTest.java
@@ -7,12 +7,11 @@ import javax.net.ssl.SSLContext;
 import java.security.NoSuchAlgorithmException;
 import java.util.Iterator;
 import java.util.Map;
-import java.util.concurrent.TimeUnit;
 
 import static org.hamcrest.CoreMatchers.equalTo;
-import static org.hamcrest.CoreMatchers.not;
 import static org.hamcrest.core.Is.is;
 import static org.hamcrest.core.IsNull.nullValue;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThat;
 
 /**
@@ -57,4 +56,18 @@ public class ConnectionParamsTest {
         assertThat(header3.getValue(), equalTo("Apple"));
     }
 
+    @Test
+    public void header_providers_are_registered() {
+        ConnectionParams.HeaderProvider dummyProvider1 = () -> "fooValue";
+        ConnectionParams.HeaderProvider dummyProvider2 = () -> "barValue";
+        ConnectionParams params = new ConnectionParams.Builder()
+                .addDynamicHeader("foo", dummyProvider1)
+                .addDynamicHeader("bar", dummyProvider2)
+                .build();
+        Map<String, ConnectionParams.HeaderProvider> headerProviders = params.getDynamicHeaders();
+        assertEquals(2, headerProviders.size());
+        assertEquals(dummyProvider1, headerProviders.get("foo"));
+        assertEquals(dummyProvider2, headerProviders.get("bar"));
+    }
+
 }
diff --git a/vespa-http-client/src/test/java/com/yahoo/vespa/http/client/core/communication/ApacheGatewayConnectionTest.java b/vespa-http-client/src/test/java/com/yahoo/vespa/http/client/core/communication/ApacheGatewayConnectionTest.java
index 48fe0461eab..456e5184f6e 100644
--- a/vespa-http-client/src/test/java/com/yahoo/vespa/http/client/core/communication/ApacheGatewayConnectionTest.java
+++ b/vespa-http-client/src/test/java/com/yahoo/vespa/http/client/core/communication/ApacheGatewayConnectionTest.java
@@ -32,11 +32,14 @@ import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.core.Is.is;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.stub;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 
@@ -244,6 +247,51 @@ public class ApacheGatewayConnectionTest {
         assertTrue(verifyContentSentLatch.await(10, TimeUnit.SECONDS));
     }
 
+    @Test
+    public void dynamic_headers_are_added_to_the_response() throws IOException, ServerResponseException, InterruptedException {
+        ConnectionParams.HeaderProvider headerProvider = mock(ConnectionParams.HeaderProvider.class);
+        when(headerProvider.getHeaderValue())
+                .thenReturn("v1")
+                .thenReturn("v2")
+                .thenReturn("v3");
+
+        ConnectionParams connectionParams = new ConnectionParams.Builder()
+                .addDynamicHeader("foo", headerProvider)
+                .build();
+
+        CountDownLatch verifyContentSentLatch = new CountDownLatch(1);
+
+        AtomicInteger counter = new AtomicInteger(1);
+        ApacheGatewayConnection.HttpClientFactory mockFactory = mockHttpClientFactory(post  -> {
+            Header[] fooHeader = post.getHeaders("foo");
+            assertEquals(1, fooHeader.length);
+            assertEquals("foo", fooHeader[0].getName());
+            assertEquals("v" + counter.getAndIncrement(), fooHeader[0].getValue());
+            verifyContentSentLatch.countDown();
+            return httpResponse("clientId", "3");
+
+        });
+
+        ApacheGatewayConnection apacheGatewayConnection =
+            new ApacheGatewayConnection(
+                    Endpoint.create("hostname", 666, false),
+                    new FeedParams.Builder().build(),
+                    "",
+                    connectionParams,
+                    mockFactory,
+                    "clientId");
+        apacheGatewayConnection.connect();
+        apacheGatewayConnection.handshake();
+
+        List<Document> documents = new ArrayList<>();
+        documents.add(createDoc("42", "content", true));
+        apacheGatewayConnection.writeOperations(documents);
+        apacheGatewayConnection.writeOperations(documents);
+        assertTrue(verifyContentSentLatch.await(10, TimeUnit.SECONDS));
+
+        verify(headerProvider, times(3)).getHeaderValue(); // 1x connect(), 2x writeOperations()
+    }
+
     private static ApacheGatewayConnection.HttpClientFactory mockHttpClientFactory(HttpExecuteMock httpExecuteMock) throws IOException {
         ApacheGatewayConnection.HttpClientFactory mockFactory =
                 mock(ApacheGatewayConnection.HttpClientFactory.class);