remove changes to sea test

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

Author: varun-edachali-dbx

tests/unit/test_sea_backend.py

@@ -6,7 +6,7 @@
 """
 
 import pytest
-from unittest.mock import Mock, patch, MagicMock
+from unittest.mock import patch, MagicMock, Mock
 
 from databricks.sql.backend.sea.backend import (
     SeaDatabricksClient,
@@ -216,17 +216,18 @@ def test_command_execution_sync(
             },
             "result": {"data": [["value1"]]},
         }
+        mock_http_client._make_request.return_value = execute_response
 
         with patch.object(
             sea_client, "get_execution_result", return_value="mock_result_set"
         ) as mock_get_result:
             result = sea_client.execute_command(
                 operation="SELECT 1",
-                session_id=session_id,
+                session_id=sea_session_id,
                 max_rows=100,
                 max_bytes=1000,
                 lz4_compression=False,
-                cursor=cursor,
+                cursor=mock_cursor,
                 use_cloud_fetch=False,
                 parameters=[],
                 async_op=False,
@@ -275,7 +276,7 @@ def test_command_execution_async(
             max_rows=100,
             max_bytes=1000,
             lz4_compression=False,
-            cursor=cursor,
+            cursor=mock_cursor,
             use_cloud_fetch=False,
             parameters=[],
             async_op=True,

tests/unit/test_sea_result_set.py

@@ -1,17 +1,19 @@
 """
 Tests for the SeaResultSet class.
+
+This module contains tests for the SeaResultSet class, which implements
+the result set functionality for the SEA (Statement Execution API) backend.
 """
 
 import pytest
-import unittest 
 from unittest.mock import patch, MagicMock, Mock
 
 from databricks.sql.result_set import SeaResultSet
 from databricks.sql.backend.types import CommandId, CommandState, BackendType
 
 
-class TestSeaResultSet(unittest.TestCase):
-    """Tests for the SeaResultSet class."""
+class TestSeaResultSet:
+    """Test suite for the SeaResultSet class."""
 
     @pytest.fixture
     def mock_connection(self):
@@ -40,130 +42,81 @@ def execute_response(self):
         mock_response.is_staging_operation = False
         return mock_response
 
-    def test_init_with_inline_data(self):
-        """Test initialization with inline data."""
-        # Create mock result data and manifest
-        from databricks.sql.backend.sea.models.base import ResultData, ResultManifest
-
-        result_data = ResultData(
-            data=[[1, "Alice"], [2, "Bob"], [3, "Charlie"]], external_links=None
-        )
-        manifest = ResultManifest(
-            format="JSON_ARRAY",
-            schema={},
-            total_row_count=3,
-            total_byte_count=0,
-            total_chunk_count=1,
-            truncated=False,
-            chunks=None,
-            result_compression=None,
-        )
-
+    def test_init_with_execute_response(
+        self, mock_connection, mock_sea_client, execute_response
+    ):
+        """Test initializing SeaResultSet with an execute response."""
         result_set = SeaResultSet(
-            connection=self.mock_connection,
-            execute_response=self.mock_execute_response_inline,
-            sea_client=self.mock_backend,
+            connection=mock_connection,
+            execute_response=execute_response,
+            sea_client=mock_sea_client,
             buffer_size_bytes=1000,
             arraysize=100,
-            result_data=result_data,
-            manifest=manifest,
         )
 
-        # Check properties
-        self.assertEqual(result_set.backend, self.mock_backend)
-        self.assertEqual(result_set.buffer_size_bytes, 1000)
-        self.assertEqual(result_set.arraysize, 100)
-
-        # Check statement ID
-        self.assertEqual(result_set.statement_id, "test-statement-id")
-
-        # Check status
-        self.assertEqual(result_set.status, CommandState.SUCCEEDED)
-
-        # Check description
-        self.assertEqual(result_set.description, self.sample_description)
-
-        # Check results queue
-        self.assertTrue(isinstance(result_set.results, JsonQueue))
-
-    def test_init_without_result_data(self):
-        """Test initialization without result data."""
-        # Create a result set without providing result_data
+        # Verify basic properties
+        assert result_set.command_id == execute_response.command_id
+        assert result_set.status == CommandState.SUCCEEDED
+        assert result_set.connection == mock_connection
+        assert result_set.backend == mock_sea_client
+        assert result_set.buffer_size_bytes == 1000
+        assert result_set.arraysize == 100
+        assert result_set.description == execute_response.description
+
+    def test_close(self, mock_connection, mock_sea_client, execute_response):
+        """Test closing a result set."""
         result_set = SeaResultSet(
-            connection=self.mock_connection,
-            execute_response=self.mock_execute_response_inline,
-            sea_client=self.mock_backend,
+            connection=mock_connection,
+            execute_response=execute_response,
+            sea_client=mock_sea_client,
             buffer_size_bytes=1000,
             arraysize=100,
         )
 
-        # Check properties
-        self.assertEqual(result_set.backend, self.mock_backend)
-        self.assertEqual(result_set.statement_id, "test-statement-id")
-        self.assertEqual(result_set.status, CommandState.SUCCEEDED)
-        self.assertEqual(result_set.description, self.sample_description)
-        self.assertTrue(isinstance(result_set.results, JsonQueue))
-
-        # Verify that the results queue is empty
-        self.assertEqual(result_set.results.data_array, [])
-
-    def test_init_with_error(self):
-        """Test initialization with error response."""
-        result_set = SeaResultSet(
-            connection=self.mock_connection,
-            execute_response=self.mock_execute_response_error,
-            sea_client=self.mock_backend,
-        )
+        # Close the result set
+        result_set.close()
 
-        # Check status
-        self.assertEqual(result_set.status, CommandState.FAILED)
-
-        # Check that description is None
-        self.assertIsNone(result_set.description)
-
-    def test_close(self):
-        """Test closing the result set."""
-        # Setup
-        from databricks.sql.backend.sea.models.base import ResultData, ResultManifest
-
-        result_data = ResultData(data=[[1, "Alice"]], external_links=None)
-        manifest = ResultManifest(
-            format="JSON_ARRAY",
-            schema={},
-            total_row_count=1,
-            total_byte_count=0,
-            total_chunk_count=1,
-            truncated=False,
-            chunks=None,
-            result_compression=None,
-        )
+        # Verify the backend's close_command was called
+        mock_sea_client.close_command.assert_called_once_with(result_set.command_id)
+        assert result_set.has_been_closed_server_side is True
+        assert result_set.status == CommandState.CLOSED
 
+    def test_close_when_already_closed_server_side(
+        self, mock_connection, mock_sea_client, execute_response
+    ):
+        """Test closing a result set that has already been closed server-side."""
         result_set = SeaResultSet(
-            connection=self.mock_connection,
-            execute_response=self.mock_execute_response_inline,
-            sea_client=self.mock_backend,
-            result_data=result_data,
-            manifest=manifest,
+            connection=mock_connection,
+            execute_response=execute_response,
+            sea_client=mock_sea_client,
+            buffer_size_bytes=1000,
+            arraysize=100,
         )
+        result_set.has_been_closed_server_side = True
 
-        # Mock the backend's close_command method
-        self.mock_backend.close_command = MagicMock()
-
-        # Execute
+        # Close the result set
         result_set.close()
 
-        # Verify
-        self.mock_backend.close_command.assert_called_once_with(self.mock_command_id)
+        # Verify the backend's close_command was NOT called
+        mock_sea_client.close_command.assert_not_called()
+        assert result_set.has_been_closed_server_side is True
+        assert result_set.status == CommandState.CLOSED
 
-    def test_is_staging_operation(self):
-        """Test is_staging_operation property."""
+    def test_close_when_connection_closed(
+        self, mock_connection, mock_sea_client, execute_response
+    ):
+        """Test closing a result set when the connection is closed."""
+        mock_connection.open = False
         result_set = SeaResultSet(
-            connection=self.mock_connection,
-            execute_response=self.mock_execute_response_inline,
-            sea_client=self.mock_backend,
+            connection=mock_connection,
+            execute_response=execute_response,
+            sea_client=mock_sea_client,
+            buffer_size_bytes=1000,
+            arraysize=100,
         )
 
-        self.assertFalse(result_set.is_staging_operation)
+        # Close the result set
+        result_set.close()
 
         # Verify the backend's close_command was NOT called
         mock_sea_client.close_command.assert_not_called()