1 files changed, 7 insertions, 4 deletions
diff --git a/core/src/test/java/org/elasticsearch/transport/netty/NettyTransportIT.java b/core/src/test/java/org/elasticsearch/transport/netty/NettyTransportIT.java
index 55f9bc49df..f7b8ede602 100644
--- a/core/src/test/java/org/elasticsearch/transport/netty/NettyTransportIT.java
+++ b/core/src/test/java/org/elasticsearch/transport/netty/NettyTransportIT.java
@@ -22,6 +22,7 @@ import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.Version;
 import org.elasticsearch.action.admin.cluster.health.ClusterHealthResponse;
 import org.elasticsearch.client.Client;
+import org.elasticsearch.client.transport.TransportClient;
 import org.elasticsearch.cluster.health.ClusterHealthStatus;
 import org.elasticsearch.common.component.Lifecycle;
 import org.elasticsearch.common.inject.Inject;
@@ -34,6 +35,7 @@ import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.transport.InetSocketTransportAddress;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.concurrent.AbstractRunnable;
+import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.test.ESIntegTestCase;
 import org.elasticsearch.test.ESIntegTestCase.ClusterScope;
@@ -49,6 +51,7 @@ import org.jboss.netty.channel.ChannelPipelineFactory;
 import java.io.IOException;
 import java.net.InetSocketAddress;
 import java.util.Collection;
+import java.util.Collections;
 
 import static org.elasticsearch.common.settings.Settings.settingsBuilder;
 import static org.hamcrest.Matchers.containsString;
@@ -78,9 +81,8 @@ public class NettyTransportIT extends ESIntegTestCase {
         Client transportClient = internalCluster().transportClient();
         ClusterHealthResponse clusterIndexHealths = transportClient.admin().cluster().prepareHealth().get();
         assertThat(clusterIndexHealths.getStatus(), is(ClusterHealthStatus.GREEN));
-
         try {
-            transportClient.admin().cluster().prepareHealth().putHeader("ERROR", "MY MESSAGE").get();
+            transportClient.filterWithHeader(Collections.singletonMap("ERROR", "MY MESSAGE")).admin().cluster().prepareHealth().get();
             fail("Expected exception, but didnt happen");
         } catch (ElasticsearchException e) {
             assertThat(e.getMessage(), containsString("MY MESSAGE"));
@@ -141,8 +143,9 @@ public class NettyTransportIT extends ESIntegTestCase {
                             final TransportRequest request = reg.newRequest();
                             request.remoteAddress(new InetSocketTransportAddress((InetSocketAddress) channel.getRemoteAddress()));
                             request.readFrom(buffer);
-                            if (request.hasHeader("ERROR")) {
-                                throw new ElasticsearchException((String) request.getHeader("ERROR"));
+                            String error = threadPool.getThreadContext().getHeader("ERROR");
+                            if (error != null) {
+                                throw new ElasticsearchException(error);
                             }
                             if (reg.getExecutor() == ThreadPool.Names.SAME) {
                                 //noinspection unchecked