Browse Source

[Test] Allow trace.id in default thread context (#76089)

A new trace.id header is added by #74210. It is handled almost the same
way as x-opaque-id. Specifically, it gets passed into default thread
context. This means the existing assertion should expect it in addition
to x-opaque-id.
Yang Wang 4 years ago
parent
commit
bad4e88278

+ 1 - 0
modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpClient.java

@@ -87,6 +87,7 @@ class Netty4HttpClient implements Closeable {
             final HttpRequest httpRequest = new DefaultFullHttpRequest(HTTP_1_1, HttpMethod.GET, uris[i]);
             httpRequest.headers().add(HOST, "localhost");
             httpRequest.headers().add("X-Opaque-ID", String.valueOf(i));
+            httpRequest.headers().add("traceparent", "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01");
             requests.add(httpRequest);
         }
         return sendRequests(remoteAddress, requests);

+ 4 - 1
server/src/main/java/org/elasticsearch/transport/Transports.java

@@ -13,6 +13,7 @@ import org.elasticsearch.http.HttpServerTransport;
 import org.elasticsearch.tasks.Task;
 
 import java.util.Arrays;
+import java.util.Set;
 
 public enum Transports {
     ;
@@ -52,7 +53,9 @@ public enum Transports {
 
     public static boolean assertDefaultThreadContext(ThreadContext threadContext) {
         assert threadContext.getRequestHeadersOnly().isEmpty() ||
-            threadContext.getRequestHeadersOnly().size() == 1 && threadContext.getRequestHeadersOnly().containsKey(Task.X_OPAQUE_ID) :
+            threadContext.getRequestHeadersOnly().keySet().equals(Set.of(Task.X_OPAQUE_ID)) ||
+            threadContext.getRequestHeadersOnly().keySet().equals(Set.of(Task.TRACE_ID)) ||
+            threadContext.getRequestHeadersOnly().keySet().equals(Set.of(Task.X_OPAQUE_ID, Task.TRACE_ID)) :
             "expected empty context but was " + threadContext.getRequestHeadersOnly() + " on " + Thread.currentThread().getName();
         return true;
     }