fix more tests

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

Author: varun-edachali-dbx

tests/unit/test_client.py

@@ -186,15 +186,17 @@ def test_arraysize_buffer_size_passthrough(
 
     def test_closing_result_set_with_closed_connection_soft_closes_commands(self):
         mock_connection = Mock()
-        mock_backend = Mock()
+        mock_backend = Mock(spec=ThriftDatabricksClient)
         mock_results = Mock()
         mock_backend.fetch_results.return_value = (Mock(), False)
-        # Ensure isinstance check passes
+
+        # Ensure connection appears closed
+        type(mock_connection).open = PropertyMock(return_value=False)
+        # Ensure isinstance check passes if needed
         mock_backend.__class__ = ThriftDatabricksClient
 
         # Setup session mock on the mock_connection
         mock_session = Mock()
-        mock_session.open = False
         mock_session.backend = mock_backend
         type(mock_connection).session = PropertyMock(return_value=mock_session)
 
@@ -204,11 +206,6 @@ def test_closing_result_set_with_closed_connection_soft_closes_commands(self):
         )
         result_set.results = mock_results
 
-        # Setup session mock on the mock_connection
-        mock_session = Mock()
-        mock_session.open = False
-        type(mock_connection).session = PropertyMock(return_value=mock_session)
-
         result_set.close()
 
         self.assertFalse(mock_backend.close_command.called)

tests/unit/test_thrift_backend.py

@@ -889,12 +889,9 @@ def test_handle_execute_response_can_handle_without_direct_results(
                     auth_provider=AuthProvider(),
                     ssl_options=SSLOptions(),
                 )
-                (
-                    execute_response,
-                    _,
-                    _, 
-                    _
-                ) = thrift_backend._handle_execute_response(execute_resp, Mock())
+                (execute_response, _, _, _) = thrift_backend._handle_execute_response(
+                    execute_resp, Mock()
+                )
                 self.assertEqual(
                     execute_response.status,
                     CommandState.SUCCEEDED,
@@ -967,9 +964,12 @@ def test_use_arrow_schema_if_available(self, tcli_service_class):
             )
         )
         thrift_backend = self._make_fake_thrift_backend()
-        execute_response, _, _, arrow_schema_bytes  = thrift_backend._handle_execute_response(
-            t_execute_resp, Mock()
-        )
+        (
+            execute_response,
+            _,
+            _,
+            arrow_schema_bytes,
+        ) = thrift_backend._handle_execute_response(t_execute_resp, Mock())
 
         self.assertEqual(arrow_schema_bytes, arrow_schema_mock)
 
@@ -1048,8 +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)
@@ -1183,7 +1183,12 @@ 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(), Mock(), Mock())
+        thrift_backend._handle_execute_response.return_value = (
+            Mock(),
+            Mock(),
+            Mock(),
+            Mock(),
+        )
         cursor_mock = Mock()
 
         result = thrift_backend.execute_command(
@@ -1219,7 +1224,12 @@ 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(), 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)
@@ -1252,7 +1262,12 @@ 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(), Mock(), Mock())
+        thrift_backend._handle_execute_response.return_value = (
+            Mock(),
+            Mock(),
+            Mock(),
+            Mock(),
+        )
         cursor_mock = Mock()
 
         result = thrift_backend.get_schemas(
@@ -1294,7 +1309,12 @@ 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(), Mock(), Mock())
+        thrift_backend._handle_execute_response.return_value = (
+            Mock(),
+            Mock(),
+            Mock(),
+            Mock(),
+        )
         cursor_mock = Mock()
 
         result = thrift_backend.get_tables(
@@ -1340,7 +1360,12 @@ 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(), Mock(), Mock())
+        thrift_backend._handle_execute_response.return_value = (
+            Mock(),
+            Mock(),
+            Mock(),
+            Mock(),
+        )
         cursor_mock = Mock()
 
         result = thrift_backend.get_columns(
@@ -2258,8 +2283,8 @@ def test_execute_command_sets_complex_type_fields_correctly(
         mock_handle_execute_response.return_value = (
             mock_execute_response,
             mock_arrow_schema,
-            Mock(), 
-            Mock()
+            Mock(),
+            Mock(),
         )
 
         # Iterate through each possible combination of native types (True, False and unset)