fix sea connector tests

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

examples/experimental/sea_connector_test.py

@@ -18,7 +18,6 @@
     "test_sea_sync_query",
     "test_sea_async_query",
     "test_sea_metadata",
-    "test_sea_multi_chunk",
 ]
 
 
@@ -28,12 +27,6 @@ def run_test_module(module_name: str) -> bool:
         os.path.dirname(os.path.abspath(__file__)), "tests", f"{module_name}.py"
     )
 
-    # Handle the multi-chunk test which is in the main directory
-    if module_name == "test_sea_multi_chunk":
-        module_path = os.path.join(
-            os.path.dirname(os.path.abspath(__file__)), f"{module_name}.py"
-        )
-
     # Simply run the module as a script - each module handles its own test execution
     result = subprocess.run(
         [sys.executable, module_path], capture_output=True, text=True

examples/experimental/tests/test_sea_async_query.py

@@ -77,24 +77,29 @@ def test_sea_async_query_with_cloud_fetch():
 
         logger.info("Query is no longer pending, getting results...")
         cursor.get_async_execution_result()
-        
+
         results = [cursor.fetchone()]
         results.extend(cursor.fetchmany(10))
         results.extend(cursor.fetchall())
-        logger.info(f"{len(results)} rows retrieved against 100 requested")
+        actual_row_count = len(results)
+        logger.info(
+            f"{actual_row_count} rows retrieved against {requested_row_count} requested"
+        )
 
         logger.info(
             f"Requested {requested_row_count} rows, received {actual_row_count} rows"
         )
-        
+
         # Verify total row count
         if actual_row_count != requested_row_count:
             logger.error(
                 f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}"
             )
             return False
-        
-        logger.info("PASS: Received correct number of rows with cloud fetch and all fetch methods work correctly")
+
+        logger.info(
+            "PASS: Received correct number of rows with cloud fetch and all fetch methods work correctly"
+        )
 
         # Close resources
         cursor.close()
@@ -182,20 +187,25 @@ def test_sea_async_query_without_cloud_fetch():
         results = [cursor.fetchone()]
         results.extend(cursor.fetchmany(10))
         results.extend(cursor.fetchall())
-        logger.info(f"{len(results)} rows retrieved against 100 requested")
+        actual_row_count = len(results)
+        logger.info(
+            f"{actual_row_count} rows retrieved against {requested_row_count} requested"
+        )
 
         logger.info(
             f"Requested {requested_row_count} rows, received {actual_row_count} rows"
         )
-        
+
         # Verify total row count
         if actual_row_count != requested_row_count:
             logger.error(
                 f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}"
             )
             return False
 
-        logger.info("PASS: Received correct number of rows without cloud fetch and all fetch methods work correctly")
+        logger.info(
+            "PASS: Received correct number of rows without cloud fetch and all fetch methods work correctly"
+        )
 
         # Close resources
         cursor.close()

examples/experimental/tests/test_sea_sync_query.py

@@ -6,7 +6,7 @@
 import logging
 from databricks.sql.client import Connection
 
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
 
@@ -62,10 +62,14 @@ def test_sea_sync_query_with_cloud_fetch():
         logger.info(
             f"Executing synchronous query with cloud fetch to generate {requested_row_count} rows"
         )
+        cursor.execute(query)
         results = [cursor.fetchone()]
         results.extend(cursor.fetchmany(10))
         results.extend(cursor.fetchall())
-        logger.info(f"{len(results)} rows retrieved against 100 requested")
+        actual_row_count = len(results)
+        logger.info(
+            f"{actual_row_count} rows retrieved against {requested_row_count} requested"
+        )
 
         # Close resources
         cursor.close()