test: deflake context propagation test

Change-Id: I382135f034cb5bbaa212d569b7f490568386df57

Author: igorbernstein2

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java

@@ -39,6 +39,7 @@
 import com.google.api.gax.grpc.GaxGrpcProperties;
 import com.google.api.gax.grpc.GrpcCallContext;
 import com.google.api.gax.grpc.GrpcTransportChannel;
+import com.google.api.gax.retrying.RetrySettings;
 import com.google.api.gax.rpc.FailedPreconditionException;
 import com.google.api.gax.rpc.FixedTransportChannelProvider;
 import com.google.api.gax.rpc.InstantiatingWatchdogProvider;
@@ -110,6 +111,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;
@@ -708,7 +710,14 @@ public void testCallContextPropagatedInMutationBatcher()
 
       // Override the timeout
       GrpcCallContext clientCtx =
-          GrpcCallContext.createDefault().withTimeout(Duration.ofMinutes(10));
+          GrpcCallContext.createDefault()
+                  .withTimeout(Duration.ofMinutes(10))
+                  .withRetrySettings(
+                          RetrySettings.newBuilder()
+                                  .setTotalTimeout(Duration.ofMinutes(10))
+                                  .build()
+
+                  );
 
       // Send a batch
       try (Batcher<RowMutationEntry, Void> batcher =
@@ -717,7 +726,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));
     }
@@ -738,7 +747,14 @@ public void testCallContextPropagatedInReadBatcher()
 
       // Override the timeout
       GrpcCallContext clientCtx =
-          GrpcCallContext.createDefault().withTimeout(Duration.ofMinutes(10));
+              GrpcCallContext.createDefault()
+                      .withTimeout(Duration.ofMinutes(10))
+                      .withRetrySettings(
+                              RetrySettings.newBuilder()
+                                      .setTotalTimeout(Duration.ofMinutes(10))
+                                      .build()
+
+                      );
 
       // Send a batch
       try (Batcher<ByteString, Row> batcher =
@@ -747,7 +763,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 +986,37 @@ public <ReqT, RespT> Listener<ReqT> interceptCall(
   }
 
   private static class ContextInterceptor implements ServerInterceptor {
-    final BlockingQueue<Context> contexts = Queues.newLinkedBlockingDeque();
+    final BlockingQueue<MethodContext> 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 <ReqT, RespT> Listener<ReqT> interceptCall(
         ServerCall<ReqT, RespT> serverCall,
         Metadata metadata,
         ServerCallHandler<ReqT, RespT> 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 {