test: unflake ReadAsyncTest

Author: olavloite

google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java

@@ -112,12 +112,12 @@ public void before() {
 
   @After
   public void after() {
-    spanner.close();
     mockSpanner.removeAllExecutionTimes();
+    spanner.close();
   }
 
   @Test
-  public void readAsyncPropagatesError() throws Exception {
+  public void readAsyncPropagatesError() {
     ApiFuture<Void> result;
     try (AsyncResultSet resultSet =
         client
@@ -156,6 +156,8 @@ public void emptyReadAsync() throws Exception {
                     case DONE:
                       assertThat(rs.getType()).isEqualTo(READ_TABLE_TYPE);
                       return CallbackResponse.DONE;
+                    default:
+                      fail("unknown cursor state");
                   }
                 }
               });
@@ -189,21 +191,26 @@ public void invalidDatabase() {
         SimulatedExecutionTime.stickyDatabaseNotFoundException("invalid-database"));
     mockSpanner.setBatchCreateSessionsExecutionTime(
         SimulatedExecutionTime.stickyDatabaseNotFoundException("invalid-database"));
-    mockSpanner.freeze();
-    DatabaseClient invalidClient =
-        spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, "invalid-database"));
-    ApiFuture<Struct> row =
-        invalidClient
-            .singleUse(TimestampBound.strong())
-            .readRowAsync(READ_TABLE_NAME, Key.of("k99"), READ_COLUMN_NAMES);
-    mockSpanner.unfreeze();
-    assertThrows(DatabaseNotFoundException.class, () -> get(row));
+    try {
+      mockSpanner.freeze();
+      DatabaseClient invalidClient =
+          spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, "invalid-database"));
+      ApiFuture<Struct> row =
+          invalidClient
+              .singleUse(TimestampBound.strong())
+              .readRowAsync(READ_TABLE_NAME, Key.of("k99"), READ_COLUMN_NAMES);
+      mockSpanner.unfreeze();
+      assertThrows(DatabaseNotFoundException.class, () -> get(row));
+    } finally {
+      mockSpanner.removeAllExecutionTimes();
+      mockSpanner.unfreeze();
+    }
   }
 
   @Test
-  public void tableNotFound() throws Exception {
+  public void tableNotFound() {
     mockSpanner.setStreamingReadExecutionTime(
-        SimulatedExecutionTime.ofException(
+        SimulatedExecutionTime.ofStickyException(
             Status.NOT_FOUND
                 .withDescription("Table not found: BadTableName")
                 .asRuntimeException()));
@@ -282,7 +289,7 @@ public void closeTransactionBeforeEndOfAsyncQuery() throws Exception {
     List<String> resultList = new ArrayList<>();
     do {
       results.drainTo(resultList);
-    } while (!finished.isDone() || results.size() > 0);
+    } while (!finished.isDone() || !results.isEmpty());
     assertThat(finished.get()).isTrue();
     assertThat(resultList).containsExactly("k1", "k2", "k3");
     // The session will be released back into the pool by the asynchronous result set when it has