From a972611358b1eb3e3c4ed28b5e8aa1f4c6607381 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Tue, 18 Mar 2025 16:38:32 +0100
Subject: [PATCH] test: cleanup AsyncResultSetImplTest

Cleans up some warnings and a leaked thread that was not stopped in one
of the tests. The actual flaky test was probably fixed as part of the
refactor of async queries, and is no longer reproducible.

Closes #3365
---
 .../cloud/spanner/AsyncResultSetImplTest.java | 38 +++++++++----------
 1 file changed, 18 insertions(+), 20 deletions(-)

diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java
index 0ba924ef740..74487283c1c 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java
@@ -104,7 +104,7 @@ public void tryNextNotAllowed() {
         new AsyncResultSetImpl(
             mockedProvider, mock(ResultSet.class), AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
       rs.setCallback(mock(Executor.class), mock(ReadyCallback.class));
-      IllegalStateException e = assertThrows(IllegalStateException.class, () -> rs.tryNext());
+      IllegalStateException e = assertThrows(IllegalStateException.class, rs::tryNext);
       assertThat(e.getMessage()).contains("tryNext may only be called from a DataReady callback.");
     }
   }
@@ -152,7 +152,7 @@ public void toListAsync() throws InterruptedException, ExecutionException {
   }
 
   @Test
-  public void toListAsyncPropagatesError() throws InterruptedException {
+  public void toListAsyncPropagatesError() {
     ExecutorService executor = Executors.newFixedThreadPool(1);
     ResultSet delegate = mock(ResultSet.class);
     when(delegate.next())
@@ -326,10 +326,7 @@ public void testCallbackIsNotCalledWhilePaused() throws InterruptedException, Ex
               @Override
               public Boolean answer(InvocationOnMock invocation) throws Throwable {
                 row++;
-                if (row > simulatedRows) {
-                  return false;
-                }
-                return true;
+                return row <= simulatedRows;
               }
             });
     when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class));
@@ -345,17 +342,17 @@ public Boolean answer(InvocationOnMock invocation) throws Throwable {
                 assertFalse(paused.get());
                 callbackCounter.incrementAndGet();
                 try {
-                  while (true) {
-                    switch (resultSet.tryNext()) {
-                      case OK:
-                        paused.set(true);
-                        queue.put(new Object());
-                        return CallbackResponse.PAUSE;
-                      case DONE:
-                        return CallbackResponse.DONE;
-                      case NOT_READY:
-                        return CallbackResponse.CONTINUE;
-                    }
+                  switch (resultSet.tryNext()) {
+                    case OK:
+                      paused.set(true);
+                      queue.put(new Object());
+                      return CallbackResponse.PAUSE;
+                    case DONE:
+                      return CallbackResponse.DONE;
+                    case NOT_READY:
+                      return CallbackResponse.CONTINUE;
+                    default:
+                      throw new IllegalStateException();
                   }
                 } catch (InterruptedException e) {
                   throw SpannerExceptionFactory.propagateInterrupt(e);
@@ -384,9 +381,8 @@ public Boolean answer(InvocationOnMock invocation) throws Throwable {
   }
 
   @Test
-  public void testCallbackIsNotCalledWhilePausedAndCanceled()
-      throws InterruptedException, ExecutionException {
-    Executor executor = Executors.newSingleThreadExecutor();
+  public void testCallbackIsNotCalledWhilePausedAndCanceled() {
+    ExecutorService executor = Executors.newSingleThreadExecutor();
     StreamingResultSet delegate = mock(StreamingResultSet.class);
 
     final AtomicInteger callbackCounter = new AtomicInteger();
@@ -414,6 +410,8 @@ public void testCallbackIsNotCalledWhilePausedAndCanceled()
       SpannerException exception = assertThrows(SpannerException.class, () -> get(callbackResult));
       assertEquals(ErrorCode.CANCELLED, exception.getErrorCode());
       assertEquals(1, callbackCounter.get());
+    } finally {
+      executor.shutdown();
     }
   }