fix: Ensure client closes after tests

spiceai · Oct 29, 2024 · d218720 · d218720
1 parent 68d6fb0
commit d218720
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/src/test/java/ai/spice/FlightQueryTest.java b/src/test/java/ai/spice/FlightQueryTest.java
@@ -60,6 +60,8 @@ public void testQuerySpiceCloudPlatform() throws ExecutionException, Interrupted
                 totalRows += root.getRowCount();
             }
 
+            spiceClient.close();
+
             assertEquals("Expected column count does not match", 4, columnCount);
             assertEquals("Expected row count does not match", 2000, totalRows);
 
@@ -87,6 +89,8 @@ public void testQuerySpiceOSS() throws ExecutionException, InterruptedException
                 totalRows += root.getRowCount();
             }
 
+            spiceClient.close();
+
             assertEquals("Expected column count does not match", 3, columnCount);
             assertEquals("Expected row count does not match", 10, totalRows);
 
@@ -109,6 +113,8 @@ public void testRefreshSpiceOSS() throws ExecutionException, InterruptedExceptio
                 assertTrue("Should correctly pass response message when unable to refresh table",
                         e.getMessage().contains("\"message\":"));
             }
+
+            spiceClient.close();
         } catch (Exception e) {
             fail("Should not throw exception: " + e.getMessage());
         }
@@ -147,6 +153,8 @@ public void testRefreshWithOptionsSpiceOSS() throws ExecutionException, Interrup
                 postRefreshRows += root.getRowCount();
             }
 
+            spiceClient.close();
+
             assertEquals("Expected row count does not match", 10, postRefreshRows);
         } catch (Exception e) {
             fail("Should not throw exception: " + e.getMessage());