diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java index aba814b05d..fbafe50f47 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java @@ -110,6 +110,7 @@ import io.grpc.ManagedChannelBuilder; import io.grpc.Metadata; import io.grpc.Metadata.Key; +import io.grpc.MethodDescriptor; import io.grpc.Server; import io.grpc.ServerCall; import io.grpc.ServerCall.Listener; @@ -717,7 +718,7 @@ public void testCallContextPropagatedInMutationBatcher() } // Ensure that the server got the overriden deadline - Context serverCtx = contextInterceptor.contexts.poll(); + Context serverCtx = contextInterceptor.pollContext(BigtableGrpc.getMutateRowsMethod()); assertThat(serverCtx).isNotNull(); assertThat(serverCtx.getDeadline()).isAtLeast(Deadline.after(8, TimeUnit.MINUTES)); } @@ -747,7 +748,7 @@ public void testCallContextPropagatedInReadBatcher() } // Ensure that the server got the overriden deadline - Context serverCtx = contextInterceptor.contexts.poll(); + Context serverCtx = contextInterceptor.pollContext(BigtableGrpc.getReadRowsMethod()); assertThat(serverCtx).isNotNull(); assertThat(serverCtx.getDeadline()).isAtLeast(Deadline.after(8, TimeUnit.MINUTES)); } @@ -970,16 +971,37 @@ public Listener interceptCall( } private static class ContextInterceptor implements ServerInterceptor { - final BlockingQueue contexts = Queues.newLinkedBlockingDeque(); + final BlockingQueue contexts = Queues.newLinkedBlockingDeque(); + + static class MethodContext { + final MethodDescriptor method; + final Context context; + + MethodContext(MethodDescriptor method, Context context) { + this.method = method; + this.context = context; + } + } @Override public Listener interceptCall( ServerCall serverCall, Metadata metadata, ServerCallHandler serverCallHandler) { - contexts.add(Context.current()); + contexts.add(new MethodContext(serverCall.getMethodDescriptor(), Context.current())); return serverCallHandler.startCall(serverCall, metadata); } + + Context pollContext(MethodDescriptor method) { + ContextInterceptor.MethodContext methodContext = contexts.poll(); + while (methodContext != null) { + if (method.equals(methodContext.method)) { + return methodContext.context; + } + methodContext = contexts.poll(); + } + return null; + } } private static class FakeDataService extends BigtableGrpc.BigtableImplBase {