diff --git a/src/main/java/org/elasticsearch/client/support/Headers.java b/src/main/java/org/elasticsearch/client/support/Headers.java index 7ab90322c91..719c7668892 100644 --- a/src/main/java/org/elasticsearch/client/support/Headers.java +++ b/src/main/java/org/elasticsearch/client/support/Headers.java @@ -48,7 +48,9 @@ public class Headers { public > M applyTo(M message) { for (String key : headers.names()) { - message.putHeader(key, headers.get(key)); + if (!message.hasHeader(key)) { + message.putHeader(key, headers.get(key)); + } } return message; } diff --git a/src/test/java/org/elasticsearch/client/AbstractClientHeadersTests.java b/src/test/java/org/elasticsearch/client/AbstractClientHeadersTests.java index 9a7b842a32c..1a512f5daac 100644 --- a/src/test/java/org/elasticsearch/client/AbstractClientHeadersTests.java +++ b/src/test/java/org/elasticsearch/client/AbstractClientHeadersTests.java @@ -20,6 +20,7 @@ package org.elasticsearch.client; import com.google.common.base.Throwables; +import com.google.common.collect.ImmutableMap; import org.elasticsearch.action.ActionListener; import org.elasticsearch.action.GenericAction; import org.elasticsearch.action.admin.cluster.node.shutdown.NodesShutdownAction; @@ -132,22 +133,45 @@ public abstract class AbstractClientHeadersTests extends ElasticsearchTestCase { client.admin().indices().prepareFlush().execute().addListener(new AssertingActionListener(FlushAction.NAME)); } - protected static void assertHeaders(Map headers) { + @Test + public void testOverideHeader() throws Exception { + String key1Val = randomAsciiOfLength(5); + Map expected = ImmutableMap.builder() + .put("key1", key1Val) + .put("key2", "val 2") + .build(); + + client.prepareGet("idx", "type", "id") + .putHeader("key1", key1Val) + .execute().addListener(new AssertingActionListener(GetAction.NAME, expected)); + + client.admin().cluster().prepareClusterStats() + .putHeader("key1", key1Val) + .execute().addListener(new AssertingActionListener(ClusterStatsAction.NAME, expected)); + + client.admin().indices().prepareCreate("idx") + .putHeader("key1", key1Val) + .execute().addListener(new AssertingActionListener(CreateIndexAction.NAME, expected)); + } + + protected static void assertHeaders(Map headers, Map expected) { assertThat(headers, notNullValue()); - assertThat(headers.size(), is(2)); - assertThat(headers.get("key1"), notNullValue()); - assertThat(headers.get("key1").toString(), equalTo("val1")); - assertThat(headers.get("key2"), notNullValue()); - assertThat(headers.get("key2").toString(), equalTo("val 2")); + assertThat(headers.size(), is(expected.size())); + for (Map.Entry expectedEntry : expected.entrySet()) { + assertThat(headers.get(expectedEntry.getKey()), equalTo(expectedEntry.getValue())); + } } protected static void assertHeaders(TransportMessage message) { + assertHeaders(message, HEADER_SETTINGS.getAsSettings(Headers.PREFIX).getAsStructuredMap()); + } + + protected static void assertHeaders(TransportMessage message, Map expected) { assertThat(message.getHeaders(), notNullValue()); - assertThat(message.getHeaders().size(), is(2)); - assertThat(message.getHeader("key1"), notNullValue()); - assertThat(message.getHeader("key1").toString(), equalTo("val1")); - assertThat(message.getHeader("key2"), notNullValue()); - assertThat(message.getHeader("key2").toString(), equalTo("val 2")); + assertThat(message.getHeaders().size(), is(expected.size())); + for (Map.Entry expectedEntry : expected.entrySet()) { + assertThat(message.getHeader(expectedEntry.getKey()), equalTo(expectedEntry.getValue())); + } } protected static class InternalException extends Exception { @@ -167,9 +191,15 @@ public abstract class AbstractClientHeadersTests extends ElasticsearchTestCase { protected static class AssertingActionListener implements ActionListener { private final String action; + private final Map expectedHeaders; public AssertingActionListener(String action) { + this(action, HEADER_SETTINGS.getAsSettings(Headers.PREFIX).getAsStructuredMap()); + } + + public AssertingActionListener(String action, Map expectedHeaders) { this.action = action; + this.expectedHeaders = expectedHeaders; } @Override @@ -183,7 +213,7 @@ public abstract class AbstractClientHeadersTests extends ElasticsearchTestCase { assertThat("expected action [" + action + "] to throw an internal exception", e, notNullValue()); assertThat(action, equalTo(((InternalException) e).action)); Map headers = ((InternalException) e).headers; - assertHeaders(headers); + assertHeaders(headers, expectedHeaders); } public Throwable unwrap(Throwable t, Class exceptionType) { diff --git a/src/test/java/org/elasticsearch/client/transport/TransportClientHeadersTests.java b/src/test/java/org/elasticsearch/client/transport/TransportClientHeadersTests.java index 515880c8115..46eddcc97c9 100644 --- a/src/test/java/org/elasticsearch/client/transport/TransportClientHeadersTests.java +++ b/src/test/java/org/elasticsearch/client/transport/TransportClientHeadersTests.java @@ -19,7 +19,6 @@ package org.elasticsearch.client.transport; -import org.apache.lucene.util.IOUtils; import org.elasticsearch.Version; import org.elasticsearch.action.GenericAction; import org.elasticsearch.action.admin.cluster.node.info.NodeInfo;