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,7 @@ 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,7 +258,8 @@ 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;
     scheduler.schedule(
         () -> {
           prepareFuture.set(
@@ -267,21 +268,30 @@ public void planRefreshDelayIsFactoredIntoExecuteTimeout() throws InterruptedExc
                       ByteString.copyFromUtf8("initialPlan"),
                       metadata(columnMetadata("strCol", stringType())))));
         },
-        50,
+        prepareDelayMs,
         TimeUnit.MILLISECONDS);
+
+    Duration originalAttemptTimeout = Duration.ofMillis(500);
     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 around (attemptDeadline - prepareDelay)
+    // add some padding to avoid flakiness.
+    int paddingMs = 5;
+    Instant estimatedDeadlineTs =
+        startOfCall
+            .plus(originalAttemptTimeout)
+            .minus(Duration.ofMillis(prepareDelayMs - paddingMs));
+    Deadline deadlineUpperBound =
+        Deadline.after(estimatedDeadlineTs.toEpochMilli(), TimeUnit.MILLISECONDS);
     Deadline executeDeadline = grpcCallContext.getCallOptions().getDeadline();
-    assertThat(executeDeadline.isBefore(paddedDeadlineAtStartOfCall)).isTrue();
+    assertThat(executeDeadline.isBefore(deadlineUpperBound)).isTrue();
   }
 
   @Test