diff --git a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java index 01de76298955..30a78e0475a7 100644 --- a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java +++ b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java @@ -40,6 +40,7 @@ import java.util.HashSet; import java.util.List; import java.util.UUID; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executors; import java.util.concurrent.Future; @@ -229,15 +230,27 @@ void testMultiStreamClosed_multiplexingEnabled() throws Exception { createConnectionWorkerPool( /* maxRequests= */ 3, /* maxBytes= */ 1000, java.time.Duration.ofSeconds(5)); - // Sets the sleep time to simulate requests stuck in connection. - testBigQueryWrite.setResponseSleep(Duration.ofMillis(50L)); StreamWriter writeStream1 = getTestStreamWriter(TEST_STREAM_1); StreamWriter writeStream2 = getTestStreamWriter(TEST_STREAM_2); // Try append 20 requests, at the end we should have 2 requests per connection. long appendCount = 20; + // Use a CountDownLatch to block mock server responses. This ensures all 20 requests + // remain in-flight simultaneously, forcing the ConnectionWorkerPool to scale up to + // 10 connections. Without this blocking, requests might finish early and prevent scaling. + CountDownLatch latch = new CountDownLatch(1); for (long i = 0; i < appendCount; i++) { - testBigQueryWrite.addResponse(createAppendResponse(i)); + long offset = i; + testBigQueryWrite.addResponse( + () -> { + try { + latch.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + return new FakeBigQueryWriteImpl.Response(createAppendResponse(offset)); + }); } List> futures = new ArrayList<>(); @@ -253,12 +266,18 @@ void testMultiStreamClosed_multiplexingEnabled() throws Exception { writeStream, connectionWorkerPool, new String[] {String.valueOf(i)}, i)); } + // At the end we should scale up to 10 connections. + try { + assertThat(connectionWorkerPool.getCreateConnectionCount()).isEqualTo(10); + assertThat(connectionWorkerPool.getTotalConnectionCount()).isEqualTo(10); + } finally { + // Release the latch to allow requests to complete. + latch.countDown(); + } + for (ApiFuture future : futures) { future.get(); } - // At the end we should scale up to 10 connections. - assertThat(connectionWorkerPool.getCreateConnectionCount()).isEqualTo(10); - assertThat(connectionWorkerPool.getTotalConnectionCount()).isEqualTo(10); // Start testing calling close on each stream. // When we close the first stream, only the connection that only serve stream 1 will be closed.