remove sea response as init option

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

Author: varun-edachali-dbx

tests/unit/test_sea_result_set.py

@@ -27,38 +27,6 @@ def mock_sea_client(self):
         """Create a mock SEA client."""
         return Mock()
 
-    @pytest.fixture
-    def sea_response(self):
-        """Create a sample SEA response."""
-        return {
-            "statement_id": "test-statement-123",
-            "status": {"state": "SUCCEEDED"},
-            "manifest": {
-                "format": "JSON_ARRAY",
-                "schema": {
-                    "column_count": 1,
-                    "columns": [
-                        {
-                            "name": "test_value",
-                            "type_text": "INT",
-                            "type_name": "INT",
-                            "position": 0,
-                        }
-                    ],
-                },
-                "total_chunk_count": 1,
-                "chunks": [{"chunk_index": 0, "row_offset": 0, "row_count": 1}],
-                "total_row_count": 1,
-                "truncated": False,
-            },
-            "result": {
-                "chunk_index": 0,
-                "row_offset": 0,
-                "row_count": 1,
-                "data_array": [["1"]],
-            },
-        }
-
     @pytest.fixture
     def execute_response(self):
         """Create a sample execute response."""
@@ -72,78 +40,35 @@ def execute_response(self):
             ("test_value", "INT", None, None, None, None, None)
         ]
         mock_response.is_staging_operation = False
-        mock_response.sea_response = {
-            "statement_id": "test-statement-123",
-            "status": {"state": "SUCCEEDED"},
-            "result": {"data_array": [["1"]]},
-        }
         return mock_response
 
-    def test_init_with_sea_response(
-        self, mock_connection, mock_sea_client, sea_response
-    ):
-        """Test initializing SeaResultSet with a SEA response."""
-        result_set = SeaResultSet(
-            connection=mock_connection,
-            sea_client=mock_sea_client,
-            sea_response=sea_response,
-            buffer_size_bytes=1000,
-            arraysize=100,
-        )
-
-        # Verify basic properties
-        assert result_set.statement_id == "test-statement-123"
-        assert result_set.status == CommandState.SUCCEEDED
-        assert result_set.command_id.guid == "test-statement-123"
-        assert result_set.command_id.backend_type == BackendType.SEA
-        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._response == sea_response
-
     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=mock_connection,
-            sea_client=mock_sea_client,
             execute_response=execute_response,
+            sea_client=mock_sea_client,
             buffer_size_bytes=1000,
             arraysize=100,
         )
 
         # Verify basic properties
-        assert result_set.statement_id == "test-statement-123"
+        assert result_set.command_id == execute_response.command_id
         assert result_set.status == CommandState.SUCCEEDED
-        assert result_set.command_id.guid == "test-statement-123"
-        assert result_set.command_id.backend_type == BackendType.SEA
         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._response == execute_response.sea_response
+        assert result_set.description == execute_response.description
 
-    def test_init_with_no_response(self, mock_connection, mock_sea_client):
-        """Test that initialization fails when neither response type is provided."""
-        with pytest.raises(ValueError) as excinfo:
-            SeaResultSet(
-                connection=mock_connection,
-                sea_client=mock_sea_client,
-                buffer_size_bytes=1000,
-                arraysize=100,
-            )
-        assert "Either execute_response or sea_response must be provided" in str(
-            excinfo.value
-        )
-
-    def test_close(self, mock_connection, mock_sea_client, sea_response):
+    def test_close(self, mock_connection, mock_sea_client, execute_response):
         """Test closing a result set."""
         result_set = SeaResultSet(
             connection=mock_connection,
+            execute_response=execute_response,
             sea_client=mock_sea_client,
-            sea_response=sea_response,
             buffer_size_bytes=1000,
             arraysize=100,
         )
@@ -157,13 +82,13 @@ def test_close(self, mock_connection, mock_sea_client, sea_response):
         assert result_set.status == CommandState.CLOSED
 
     def test_close_when_already_closed_server_side(
-        self, mock_connection, mock_sea_client, sea_response
+        self, mock_connection, mock_sea_client, execute_response
     ):
         """Test closing a result set that has already been closed server-side."""
         result_set = SeaResultSet(
             connection=mock_connection,
+            execute_response=execute_response,
             sea_client=mock_sea_client,
-            sea_response=sea_response,
             buffer_size_bytes=1000,
             arraysize=100,
         )
@@ -178,14 +103,14 @@ def test_close_when_already_closed_server_side(
         assert result_set.status == CommandState.CLOSED
 
     def test_close_when_connection_closed(
-        self, mock_connection, mock_sea_client, sea_response
+        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=mock_connection,
+            execute_response=execute_response,
             sea_client=mock_sea_client,
-            sea_response=sea_response,
             buffer_size_bytes=1000,
             arraysize=100,
         )
@@ -199,13 +124,13 @@ def test_close_when_connection_closed(
         assert result_set.status == CommandState.CLOSED
 
     def test_unimplemented_methods(
-        self, mock_connection, mock_sea_client, sea_response
+        self, mock_connection, mock_sea_client, execute_response
     ):
         """Test that unimplemented methods raise NotImplementedError."""
         result_set = SeaResultSet(
             connection=mock_connection,
+            execute_response=execute_response,
             sea_client=mock_sea_client,
-            sea_response=sea_response,
             buffer_size_bytes=1000,
             arraysize=100,
         )
@@ -258,18 +183,18 @@ def test_unimplemented_methods(
                 pass
 
     def test_fill_results_buffer_not_implemented(
-        self, mock_connection, mock_sea_client, sea_response
+        self, mock_connection, mock_sea_client, execute_response
     ):
         """Test that _fill_results_buffer raises NotImplementedError."""
         result_set = SeaResultSet(
             connection=mock_connection,
+            execute_response=execute_response,
             sea_client=mock_sea_client,
-            sea_response=sea_response,
             buffer_size_bytes=1000,
             arraysize=100,
         )
 
         with pytest.raises(
             NotImplementedError, match="fetchone is not implemented for SEA backend"
         ):
-            result_set._fill_results_buffer()
+            result_set._fill_results_buffer()