test: unflake RetryOnInvalidatedSessionTest (#3561)

googleapis · Dec 16, 2024 · 430143c · 430143c
1 parent 557b214
commit 430143c
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/...e-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryOnInvalidatedSessionTest.java b/...e-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryOnInvalidatedSessionTest.java
@@ -207,7 +207,7 @@ public static void stopServer() throws InterruptedException {
   public void setUp() throws InterruptedException {
     mockSpanner.reset();
     if (spanner == null
-        || spanner.getOptions().getSessionPoolOptions().isFailIfPoolExhausted()
+        || spanner.getOptions().getSessionPoolOptions().isFailIfSessionNotFound()
             != failOnInvalidatedSession) {
       if (spanner != null) {
         spanner.close();
@@ -228,8 +228,8 @@ public void setUp() throws InterruptedException {
               .build()
               .getService();
       client = spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
-      invalidateSessionPool(client, spanner.getOptions().getSessionPoolOptions().getMinSessions());
     }
+    invalidateSessionPool(client, spanner.getOptions().getSessionPoolOptions().getMinSessions());
   }
 
   private static void invalidateSessionPool(DatabaseClient client, int minSessions)
@@ -1002,6 +1002,7 @@ public void transactionManagerBatchUpdate() throws InterruptedException {
         } catch (AbortedException e) {
           transaction = assertThrowsSessionNotFoundIfShouldFail(() -> manager.resetForRetry());
           if (transaction == null) {
+            manager.close();
             break;
           }
         }