test: deflake goaway test to account for last vRpc being rejected

Change-Id: Ife308c3cd8dd693909cdd9af4f02a9ec63c16071

Author: igorbernstein2

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,24 @@ 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("Sanity check that some vRpcs succeeded prior to the goaway")
+        .that(numOk)
+        .isGreaterThan(0);
+    assertWithMessage(
+            "Sanity check possibly only the last vRpc was 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