diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallableTest.java index 3337f72ce3c4..1c28d53e127c 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallableTest.java +++ b/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 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 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); callable.call(callContext, outerObserver, context); + Instant startOfCall = Instant.now(); 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