test: deflake goaway test to account for last vRpc being rejected (#2870)

igorbernstein2 · web-flow · commit b724bfac4494 · 2026-04-09T15:38:25.000Z
Change-Id: Ife308c3cd8dd693909cdd9af4f02a9ec63c16071 Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://togithub.com/googleapis/java-bigtable/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) - [ ] Rollback plan is reviewed and LGTMed - [ ] All new data plane features have a completed end to end testing plan Fixes #<issue_number_goes_here> ☕️ If you write sample code, please follow the [samples format]( https://togithub.com/GoogleCloudPlatform/java-docs-samples/blob/main/SAMPLE_FORMAT.md).
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java
@@ -19,6 +19,7 @@
 import static com.google.cloud.bigtable.data.v2.internal.test_helpers.StatusSubject.assertThat;
 import static com.google.cloud.bigtable.data.v2.internal.test_helpers.VRpcResultSubject.assertThat;
 import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertWithMessage;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import com.google.bigtable.v2.CloseSessionRequest;
@@ -201,6 +202,8 @@ void sessionGoAwayTest() throws Exception {
 
     // Send vRPCs until after a goaway time
     Stopwatch stopwatch = Stopwatch.createStarted();
+    int numUncommittedErrors = 0;
+    int numOk = 0;
     while (stopwatch.elapsed(TimeUnit.MILLISECONDS) < goAwayDelay.toMillis()) {
       VRpc<SessionFakeScriptedRequest, SessionFakeScriptedResponse> rpc =
           session.newCall(FakeDescriptor.SCRIPTED);
@@ -209,9 +212,23 @@ void sessionGoAwayTest() throws Exception {
           SessionFakeScriptedRequest.newBuilder().setTag(0).build(),
           VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), true, tracer),
           f);
-      f.get();
+      try {
+        f.get();
+        numOk++;
+      } catch (VRpcException e) {
+        if (e.getResult().getState() == State.UNCOMMITED) {
+          numUncommittedErrors++;
+        }
+      }
     }
 
+    assertWithMessage("Ensure that some vRpcs succeeded prior to the goaway")
+        .that(numOk)
+        .isGreaterThan(0);
+    assertWithMessage("Ensure that only the last vRpc could be rejected with an uncommited error")
+        .that(numUncommittedErrors)
+        .isAtMost(1);
+
     assertThat(sessionListener.popUntil(GoAwayResponse.class)).isInstanceOf(GoAwayResponse.class);
 
     // Make sure we can't send vrpc after receiving goaway
