From 430143ccc96c4bf5bbccd717b820e301b8a6b474 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Mon, 16 Dec 2024 18:00:28 +0100
Subject: [PATCH] test: unflake RetryOnInvalidatedSessionTest (#3561)

---
 .../google/cloud/spanner/RetryOnInvalidatedSessionTest.java  | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryOnInvalidatedSessionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryOnInvalidatedSessionTest.java
index a6c3f9fa7c..f070f15421 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryOnInvalidatedSessionTest.java
+++ b/google-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;
           }
         }