test: Deflake plan refresh deadline test

Change-Id: I1605f65fe7053adee9e88387c8fcfbe60736be07

Author: jackdingilian

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

@@ -244,7 +244,8 @@ public void testIsPlanRefreshError() {
   }
 
   @Test
-  public void planRefreshDelayIsFactoredIntoExecuteTimeout() throws InterruptedException {
+  public void planRefreshDelayIsFactoredIntoExecuteTimeout()
+      throws InterruptedException, ExecutionException {
     MockServerStreamingCallable<ExecuteQueryRequest, ExecuteQueryResponse> innerCallable =
         new MockServerStreamingCallable<>();
     RequestContext requestContext = RequestContext.create("project", "instance", "profile");
@@ -258,30 +259,39 @@ public void planRefreshDelayIsFactoredIntoExecuteTimeout() throws InterruptedExc
     ExecuteQueryCallContext callContext =
         ExecuteQueryCallContext.create(preparedStatement.bind().build(), metadataFuture);
 
-    Duration originalAttemptTimeout = Duration.ofMillis(100);
+    // Simulate delay for prepare response
+    int prepareDelayMs = 50;
+    SettableApiFuture<Instant> prepareResolvedAt = SettableApiFuture.create();
     scheduler.schedule(
         () -> {
+          prepareResolvedAt.set(Instant.now());
           prepareFuture.set(
               PrepareResponse.fromProto(
                   prepareResponse(
                       ByteString.copyFromUtf8("initialPlan"),
                       metadata(columnMetadata("strCol", stringType())))));
         },
-        50,
+        prepareDelayMs,
         TimeUnit.MILLISECONDS);
+
+    Duration originalAttemptTimeout = Duration.ofMillis(5000);
     ApiCallContext context =
         GrpcCallContext.createDefault().withTimeoutDuration(originalAttemptTimeout);
-    // prepare takes 50 ms to resolve. Despite that the execute timeout should be around 100ms from
-    // now (w padding)
-    Deadline paddedDeadlineAtStartOfCall =
-        Deadline.after(originalAttemptTimeout.toMillis() + 5, TimeUnit.MILLISECONDS);
+    Instant startOfCall = Instant.now();
     callable.call(callContext, outerObserver, context);
     scheduler.shutdown();
     scheduler.awaitTermination(30, TimeUnit.SECONDS);
     GrpcCallContext grpcCallContext =
         (GrpcCallContext) innerCallable.popLastCall().getApiCallContext();
+
+    // Deadline should be at most (original deadline - prepare delay). Add 1ms to avoid flakiness.
+    Duration prepareLatency = Duration.between(startOfCall, prepareResolvedAt.get());
+    Deadline deadlineUpperBound =
+        Deadline.after(
+            originalAttemptTimeout.toMillis() - prepareLatency.toMillis() + 1,
+            TimeUnit.MILLISECONDS);
     Deadline executeDeadline = grpcCallContext.getCallOptions().getDeadline();
-    assertThat(executeDeadline.isBefore(paddedDeadlineAtStartOfCall)).isTrue();
+    assertThat(executeDeadline.isBefore(deadlineUpperBound)).isTrue();
   }
 
   @Test