From 959fda272678b4371a5d1b87d9cde889fc11bd0d Mon Sep 17 00:00:00 2001
From: Phillip LeBlanc <phillip@leblanc.tech>
Date: Thu, 26 Mar 2026 12:47:58 +0900
Subject: [PATCH] fix: Improve concurrent reset+query test to catch
 IllegalStateException

---
 src/test/java/ai/spice/ResetTest.java | 40 +++++++++++++++++++++++----
 1 file changed, 35 insertions(+), 5 deletions(-)

diff --git a/src/test/java/ai/spice/ResetTest.java b/src/test/java/ai/spice/ResetTest.java
index c92904a..e57728b 100644
--- a/src/test/java/ai/spice/ResetTest.java
+++ b/src/test/java/ai/spice/ResetTest.java
@@ -26,11 +26,13 @@ of this software and associated documentation files (the "Software"), to deal
 import java.util.List;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import org.apache.arrow.flight.FlightRuntimeException;
 import org.apache.arrow.flight.FlightStream;
 import org.apache.arrow.vector.ipc.ArrowReader;
 
@@ -294,10 +296,13 @@ public void testConcurrentResetAndQuery() throws Exception {
                         try {
                             FlightStream stream = client.query("SELECT 1");
                             stream.close();
-                        } catch (NullPointerException e) {
-                            unexpectedErrors.add(e);
                         } catch (Exception e) {
-                            // Connection errors are expected
+                            // Unwrap ExecutionException to inspect the real cause
+                            Throwable toCheck = (e instanceof ExecutionException && e.getCause() != null)
+                                    ? e.getCause() : e;
+                            if (!isExpectedTransportError(toCheck)) {
+                                unexpectedErrors.add(toCheck);
+                            }
                         }
                         Thread.sleep(10);
                     }
@@ -310,14 +315,39 @@ public void testConcurrentResetAndQuery() throws Exception {
             assertTrue("Concurrent reset+query should complete within 30s",
                     executor.awaitTermination(30, TimeUnit.SECONDS));
 
-            assertTrue("Should not get NullPointerException during concurrent reset+query: " + unexpectedErrors,
-                    unexpectedErrors.isEmpty());
+            assertTrue("Unexpected errors during concurrent reset+query: "
+                    + unexpectedErrors, unexpectedErrors.isEmpty());
         } finally {
             executor.shutdownNow();
             client.close();
         }
     }
 
+    /**
+     * Returns true if the throwable is an expected transport/connection error
+     * that can occur when no server is running or during a mid-reset race.
+     * Anything else (NPE, IllegalStateException, etc.) is unexpected.
+     */
+    private static boolean isExpectedTransportError(Throwable t) {
+        if (t instanceof FlightRuntimeException) {
+            return true; // gRPC status errors (UNAVAILABLE, etc.)
+        }
+        if (t instanceof java.net.ConnectException
+                || t instanceof java.io.IOException) {
+            return true; // connection refused, broken pipe, etc.
+        }
+        String msg = t.getMessage();
+        if (msg != null) {
+            String lower = msg.toLowerCase();
+            if (lower.contains("unavailable")
+                    || lower.contains("connection refused")
+                    || lower.contains("failed to execute")) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     // ==================== Construction / DNS / Keep-alive ====================
 
     /**