DATAES-801 - Implement callback to enable adding custom headers in the REST HTTP request.

Original PR: #442

src/test/java/org/springframework/data/elasticsearch/client/RestClientsTest.java

@@ -4,44 +4,221 @@ import static com.github.tomakehurst.wiremock.client.WireMock.*;
 import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.options;
 
 import java.io.IOException;
+import java.util.Arrays;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
+import java.util.stream.Stream;
 
 import org.elasticsearch.client.RequestOptions;
 import org.elasticsearch.client.RestHighLevelClient;
-import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.MethodSource;
+import org.springframework.data.elasticsearch.client.reactive.ReactiveElasticsearchClient;
+import org.springframework.data.elasticsearch.client.reactive.ReactiveRestClients;
+import org.springframework.http.HttpHeaders;
 
 import com.github.tomakehurst.wiremock.WireMockServer;
 import com.github.tomakehurst.wiremock.client.WireMock;
+import com.github.tomakehurst.wiremock.matching.AnythingPattern;
+import com.github.tomakehurst.wiremock.matching.EqualToPattern;
 
 /**
  * @author Peter-Josef Meisch
  */
 public class RestClientsTest {
 
-	@Test // DATAES-700
-	void shouldUseConfiguredProxy() throws IOException {
+	@ParameterizedTest // DATAES-700
+	@MethodSource("clientUnderTestFactorySource")
+	@DisplayName("should use configured proxy")
+	void shouldUseConfiguredProxy(ClientUnderTestFactory clientUnderTestFactory) throws IOException {
 
-		WireMockServer wireMockServer = new WireMockServer(options() //
-				.dynamicPort() //
-				.usingFilesUnderDirectory("src/test/resources/wiremock-mappings")); // needed, otherwise Wiremock goes to
-																					// test/resources/mappings
-		wireMockServer.start();
-		try {
-			WireMock.configureFor(wireMockServer.port());
+		if (clientUnderTestFactory instanceof ReactiveElasticsearchClientUnderTestFactory) {
+			// although the reactive code is using the proxy for every call - tested with an intercepting
+			// proxy - somehow in this test wiremock fails to register this. So we skip it here
+			//
+			return;
+		}
+
+		wireMockServer(server -> {
+
+			WireMock.configureFor(server.port());
+
+			stubFor(head(urlEqualTo("/")).willReturn(aResponse() //
+					.withHeader("Content-Type", "application/json; charset=UTF-8")));
 
 			ClientConfigurationBuilder configurationBuilder = new ClientConfigurationBuilder();
 			ClientConfiguration clientConfiguration = configurationBuilder //
-					.connectedTo("localhost:9200")//
-					.withProxy("localhost:" + wireMockServer.port()) //
+					.connectedTo("localhost:4711")//
+					.withProxy("localhost:" + server.port()) //
 					.build();
+			ClientUnderTest clientUnderTest = clientUnderTestFactory.create(clientConfiguration);
 
-			RestHighLevelClient restClient = RestClients.create(clientConfiguration).rest();
-			restClient.ping(RequestOptions.DEFAULT);
+			clientUnderTest.ping();
 
 			verify(headRequestedFor(urlEqualTo("/")));
+		});
+	}
 
+	@ParameterizedTest // DATAES-801
+	@MethodSource("clientUnderTestFactorySource")
+	@DisplayName("should set all required headers")
+	void shouldSetAllRequiredHeaders(ClientUnderTestFactory clientUnderTestFactory) {
+		wireMockServer(server -> {
+
+			WireMock.configureFor(server.port());
+
+			stubFor(head(urlEqualTo("/")).willReturn(aResponse() //
+					.withHeader("Content-Type", "application/json; charset=UTF-8")));
+
+			HttpHeaders defaultHeaders = new HttpHeaders();
+			defaultHeaders.addAll("def1", Arrays.asList("def1-1", "def1-2"));
+			defaultHeaders.add("def2", "def2-1");
+
+			AtomicInteger supplierCount = new AtomicInteger(1);
+
+			ClientConfigurationBuilder configurationBuilder = new ClientConfigurationBuilder();
+			ClientConfiguration clientConfiguration = configurationBuilder //
+					.connectedTo("localhost:" + server.port()) //
+					.withBasicAuth("user", "password") //
+					.withDefaultHeaders(defaultHeaders) //
+					.withHeaders(() -> {
+						HttpHeaders httpHeaders = new HttpHeaders();
+						httpHeaders.add("supplied", "val0");
+						httpHeaders.add("supplied", "val" + supplierCount.getAndIncrement());
+						return httpHeaders;
+					}).build();
+
+			ClientUnderTest clientUnderTest = clientUnderTestFactory.create(clientConfiguration);
+
+			// do several calls to check that the headerSupplier provided values are set
+			for (int i = 1; i <= 3; i++) {
+				clientUnderTest.ping();
+
+				verify(headRequestedFor(urlEqualTo("/")).withHeader("Authorization", new AnythingPattern()) //
+						.withHeader("def1", new EqualToPattern("def1-1")) //
+						.withHeader("def1", new EqualToPattern("def1-2")) //
+						.withHeader("def2", new EqualToPattern("def2-1")) //
+						.withHeader("supplied", new EqualToPattern("val0")) //
+						.withHeader("supplied", new EqualToPattern("val" + i)) //
+				);
+			}
+		});
+	}
+
+	/**
+	 * Consumer extension that catches checked exceptions and wraps them in a RuntimeException.
+	 */
+	@FunctionalInterface
+	interface WiremockConsumer extends Consumer<WireMockServer> {
+		@Override
+		default void accept(WireMockServer wiremockConsumer) {
+			try {
+				acceptThrows(wiremockConsumer);
+			} catch (final Exception e) {
+				throw new RuntimeException(e);
+			}
+		}
+
+		void acceptThrows(WireMockServer wiremockConsumer) throws Exception;
+	}
+
+	/**
+	 * starts a Wiremock server and calls consumer with the server as argument. Stops the server after consumer execution.
+	 * 
+	 * @param consumer the consumer
+	 */
+	private void wireMockServer(WiremockConsumer consumer) {
+		WireMockServer wireMockServer = new WireMockServer(options() //
+				.dynamicPort() //
+				.usingFilesUnderDirectory("src/test/resources/wiremock-mappings")); // needed, otherwise Wiremock goes to
+		// test/resources/mappings
+		try {
+			wireMockServer.start();
+			consumer.accept(wireMockServer);
 		} finally {
 			wireMockServer.shutdown();
 		}
 	}
 
+	/**
+	 * The client to be tested. Abstraction to be able to test reactive and non-reactive clients.
+	 */
+	interface ClientUnderTest {
+		/**
+		 * Pings the configured server.
+		 * 
+		 * @return
+		 */
+		boolean ping() throws Exception;
+	}
+
+	/**
+	 * base class to create {@link ClientUnderTest} implementations.
+	 */
+	static abstract class ClientUnderTestFactory {
+		abstract ClientUnderTest create(ClientConfiguration clientConfiguration);
+
+		@Override
+		public String toString() {
+			return getDisplayName();
+		}
+
+		protected abstract String getDisplayName();
+	}
+
+	/**
+	 * {@link ClientUnderTestFactory} implementation for the Standard {@link RestHighLevelClient}.
+	 */
+	static class RestClientUnderTestFactory extends ClientUnderTestFactory {
+
+		@Override
+		protected String getDisplayName() {
+			return "RestHighLevelClient";
+		}
+
+		@Override
+		ClientUnderTest create(ClientConfiguration clientConfiguration) {
+			RestHighLevelClient client = RestClients.create(clientConfiguration).rest();
+			return new ClientUnderTest() {
+
+				@Override
+				public boolean ping() throws Exception {
+					return client.ping(RequestOptions.DEFAULT);
+				}
+			};
+		}
+
+	}
+
+	/**
+	 * {@link ClientUnderTestFactory} implementation for the {@link ReactiveElasticsearchClient}.
+	 */
+	static class ReactiveElasticsearchClientUnderTestFactory extends ClientUnderTestFactory {
+
+		@Override
+		protected String getDisplayName() {
+			return "ReactiveElasticsearchClient";
+		}
+
+		@Override
+		ClientUnderTest create(ClientConfiguration clientConfiguration) {
+			ReactiveElasticsearchClient client = ReactiveRestClients.create(clientConfiguration);
+			return new ClientUnderTest() {
+				@Override
+				public boolean ping() throws Exception {
+					return client.ping().block();
+				}
+			};
+		}
+	}
+
+	/**
+	 * Provides the factories to use in the parameterized tests
+	 * 
+	 * @return stream of factories
+	 */
+	static Stream<ClientUnderTestFactory> clientUnderTestFactorySource() {
+		return Stream.of(new RestClientUnderTestFactory(), new ReactiveElasticsearchClientUnderTestFactory());
+	}
 }

src/test/java/org/springframework/data/elasticsearch/client/reactive/ReactiveMockClientTestsUtils.java

@@ -83,10 +83,10 @@ public class ReactiveMockClientTestsUtils {
 
 		if (hosts.length == 1) {
 
-			delegate = new SingleNodeHostProvider(clientProvider, getInetSocketAddress(hosts[0])) {};
+			delegate = new SingleNodeHostProvider(clientProvider, HttpHeaders::new, getInetSocketAddress(hosts[0])) {};
 		} else {
 
-			delegate = new MultiNodeHostProvider(clientProvider, Arrays.stream(hosts)
+			delegate = new MultiNodeHostProvider(clientProvider,HttpHeaders::new, Arrays.stream(hosts)
 					.map(ReactiveMockClientTestsUtils::getInetSocketAddress).toArray(InetSocketAddress[]::new)) {};
 		}