fix some tests

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

Author: varun-edachali-dbx

tests/unit/test_client.py

@@ -108,7 +108,6 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class):
                 mock_execute_response.status = (
                     CommandState.SUCCEEDED if not closed else CommandState.CLOSED
                 )
-                mock_execute_response.has_been_closed_server_side = closed
                 mock_execute_response.is_staging_operation = False
                 mock_execute_response.description = []
 
@@ -127,6 +126,7 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class):
                 real_result_set = ThriftResultSet(
                     connection=connection,
                     execute_response=mock_execute_response,
+                    has_been_closed_server_side=closed,
                 )
 
                 # Mock execute_command to return our real result set

tests/unit/test_fetches.py

@@ -63,12 +63,12 @@ def make_dummy_result_set_from_initial_results(initial_results):
             execute_response=ExecuteResponse(
                 command_id=None,
                 status=None,
-                has_been_closed_server_side=True,
                 description=description,
                 lz4_compressed=True,
                 is_staging_operation=False,
             ),
             t_row_set=None,
+            has_been_closed_server_side=True,
         )
         return rs
 
@@ -117,11 +117,11 @@ def fetch_results(
             execute_response=ExecuteResponse(
                 command_id=None,
                 status=None,
-                has_been_closed_server_side=False,
                 description=description,
                 lz4_compressed=True,
                 is_staging_operation=False,
             ),
+            has_been_closed_server_side=False,
         )
         return rs
 

tests/unit/test_sea_result_set.py

@@ -141,29 +141,6 @@ def test_close(self, mock_connection, execute_response):
         mock_connection.session.backend.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, execute_response
-    ):
-        """Test closing a result set that has already been closed server-side."""
-        result_set = SeaResultSet(
-            connection=mock_connection,
-            execute_response=execute_response,
-            result_data=ResultData(data=[]),
-            manifest=self._create_empty_manifest(ResultFormat.JSON_ARRAY),
-            buffer_size_bytes=1000,
-            arraysize=100,
-        )
-        result_set.has_been_closed_server_side = True
-
-        # Close the result set
-        result_set.close()
-
-        # Verify the backend's close_command was NOT called
-        mock_connection.session.backend.close_command.assert_not_called()
-        assert result_set.has_been_closed_server_side is True
         assert result_set.status == CommandState.CLOSED
 
     def test_close_when_connection_closed(self, mock_connection, execute_response):
@@ -183,7 +160,6 @@ def test_close_when_connection_closed(self, mock_connection, execute_response):
 
         # Verify the backend's close_command was NOT called
         mock_connection.session.backend.close_command.assert_not_called()
-        assert result_set.has_been_closed_server_side is True
         assert result_set.status == CommandState.CLOSED
 
     def test_init_with_result_data(self, result_set_with_data, sample_data):

tests/unit/test_thrift_backend.py

@@ -649,7 +649,7 @@ def test_handle_execute_response_sets_compression_in_direct_results(
                 ssl_options=SSLOptions(),
             )
 
-            execute_response, _ = thrift_backend._handle_execute_response(
+            execute_response, _, _, _ = thrift_backend._handle_execute_response(
                 t_execute_resp, Mock()
             )
             self.assertEqual(execute_response.lz4_compressed, lz4Compressed)
@@ -892,6 +892,8 @@ def test_handle_execute_response_can_handle_without_direct_results(
                 (
                     execute_response,
                     _,
+                    _, 
+                    _
                 ) = thrift_backend._handle_execute_response(execute_resp, Mock())
                 self.assertEqual(
                     execute_response.status,
@@ -965,11 +967,11 @@ def test_use_arrow_schema_if_available(self, tcli_service_class):
             )
         )
         thrift_backend = self._make_fake_thrift_backend()
-        execute_response, _ = thrift_backend._handle_execute_response(
+        execute_response, _, _, arrow_schema_bytes  = thrift_backend._handle_execute_response(
             t_execute_resp, Mock()
         )
 
-        self.assertEqual(execute_response.arrow_schema_bytes, arrow_schema_mock)
+        self.assertEqual(arrow_schema_bytes, arrow_schema_mock)
 
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_fall_back_to_hive_schema_if_no_arrow_schema(self, tcli_service_class):
@@ -997,7 +999,7 @@ def test_fall_back_to_hive_schema_if_no_arrow_schema(self, tcli_service_class):
             )
         )
         thrift_backend = self._make_fake_thrift_backend()
-        _, _ = thrift_backend._handle_execute_response(t_execute_resp, Mock())
+        _, _, _, _ = thrift_backend._handle_execute_response(t_execute_resp, Mock())
 
         self.assertEqual(
             hive_schema_mock,
@@ -1046,6 +1048,8 @@ def test_handle_execute_response_reads_has_more_rows_in_direct_results(
                 (
                     execute_response,
                     has_more_rows_result,
+                    _, 
+                    _
                 ) = thrift_backend._handle_execute_response(execute_resp, Mock())
 
                 self.assertEqual(is_direct_results, has_more_rows_result)
@@ -1179,7 +1183,7 @@ def test_execute_statement_calls_client_and_handle_execute_response(
             ssl_options=SSLOptions(),
         )
         thrift_backend._handle_execute_response = Mock()
-        thrift_backend._handle_execute_response.return_value = (Mock(), Mock())
+        thrift_backend._handle_execute_response.return_value = (Mock(), Mock(), Mock(), Mock())
         cursor_mock = Mock()
 
         result = thrift_backend.execute_command(
@@ -1215,7 +1219,7 @@ def test_get_catalogs_calls_client_and_handle_execute_response(
             ssl_options=SSLOptions(),
         )
         thrift_backend._handle_execute_response = Mock()
-        thrift_backend._handle_execute_response.return_value = (Mock(), Mock())
+        thrift_backend._handle_execute_response.return_value = (Mock(), Mock(), Mock(), Mock())
         cursor_mock = Mock()
 
         result = thrift_backend.get_catalogs(Mock(), 100, 200, cursor_mock)
@@ -1248,7 +1252,7 @@ def test_get_schemas_calls_client_and_handle_execute_response(
             ssl_options=SSLOptions(),
         )
         thrift_backend._handle_execute_response = Mock()
-        thrift_backend._handle_execute_response.return_value = (Mock(), Mock())
+        thrift_backend._handle_execute_response.return_value = (Mock(), Mock(), Mock(), Mock())
         cursor_mock = Mock()
 
         result = thrift_backend.get_schemas(
@@ -1290,7 +1294,7 @@ def test_get_tables_calls_client_and_handle_execute_response(
             ssl_options=SSLOptions(),
         )
         thrift_backend._handle_execute_response = Mock()
-        thrift_backend._handle_execute_response.return_value = (Mock(), Mock())
+        thrift_backend._handle_execute_response.return_value = (Mock(), Mock(), Mock(), Mock())
         cursor_mock = Mock()
 
         result = thrift_backend.get_tables(
@@ -1336,7 +1340,7 @@ def test_get_columns_calls_client_and_handle_execute_response(
             ssl_options=SSLOptions(),
         )
         thrift_backend._handle_execute_response = Mock()
-        thrift_backend._handle_execute_response.return_value = (Mock(), Mock())
+        thrift_backend._handle_execute_response.return_value = (Mock(), Mock(), Mock(), Mock())
         cursor_mock = Mock()
 
         result = thrift_backend.get_columns(
@@ -2254,6 +2258,8 @@ def test_execute_command_sets_complex_type_fields_correctly(
         mock_handle_execute_response.return_value = (
             mock_execute_response,
             mock_arrow_schema,
+            Mock(), 
+            Mock()
         )
 
         # Iterate through each possible combination of native types (True, False and unset)