remove un-necessary changes in thrift backend tests

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

Author: varun-edachali-dbx

tests/unit/test_thrift_backend.py

@@ -623,7 +623,10 @@ def test_handle_execute_response_sets_compression_in_direct_results(
                 status=Mock(),
                 operationHandle=Mock(),
                 directResults=ttypes.TSparkDirectResults(
-                    operationStatus=op_status,
+                    operationStatus=ttypes.TGetOperationStatusResp(
+                        status=self.okay_status,
+                        operationState=ttypes.TOperationState.FINISHED_STATE,
+                    ),
                     resultSetMetadata=ttypes.TGetResultSetMetadataResp(
                         status=self.okay_status,
                         resultFormat=ttypes.TSparkRowSetType.ARROW_BASED_SET,
@@ -832,9 +835,10 @@ def test_handle_execute_response_checks_direct_results_for_error_statuses(self):
                         thrift_backend._handle_execute_response(error_resp, Mock())
                     self.assertIn("this is a bad error", str(cm.exception))
 
+    @patch("databricks.sql.backend.thrift_backend.ThriftResultSet")
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_handle_execute_response_can_handle_without_direct_results(
-        self, tcli_service_class
+        self, tcli_service_class, mock_result_set
     ):
         tcli_service_instance = tcli_service_class.return_value
 
@@ -878,10 +882,10 @@ 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,
@@ -947,8 +951,14 @@ def test_use_arrow_schema_if_available(self, tcli_service_class):
         tcli_service_instance.GetResultSetMetadata.return_value = (
             t_get_result_set_metadata_resp
         )
+        tcli_service_instance.GetOperationStatus.return_value = (
+            ttypes.TGetOperationStatusResp(
+                status=self.okay_status,
+                operationState=ttypes.TOperationState.FINISHED_STATE,
+            )
+        )
         thrift_backend = self._make_fake_thrift_backend()
-        execute_response, arrow_schema_bytes = thrift_backend._handle_execute_response(
+        execute_response, _ = thrift_backend._handle_execute_response(
             t_execute_resp, Mock()
         )
 
@@ -973,8 +983,14 @@ def test_fall_back_to_hive_schema_if_no_arrow_schema(self, tcli_service_class):
         )
 
         tcli_service_instance.GetResultSetMetadata.return_value = hive_schema_req
+        tcli_service_instance.GetOperationStatus.return_value = (
+            ttypes.TGetOperationStatusResp(
+                status=self.okay_status,
+                operationState=ttypes.TOperationState.FINISHED_STATE,
+            )
+        )
         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,
@@ -988,10 +1004,10 @@ def test_fall_back_to_hive_schema_if_no_arrow_schema(self, tcli_service_class):
     def test_handle_execute_response_reads_has_more_rows_in_direct_results(
         self, tcli_service_class, build_queue
     ):
-        for has_more_rows, resp_type in itertools.product(
+        for is_direct_results, resp_type in itertools.product(
             [True, False], self.execute_response_types
         ):
-            with self.subTest(has_more_rows=has_more_rows, resp_type=resp_type):
+            with self.subTest(is_direct_results=is_direct_results, resp_type=resp_type):
                 tcli_service_instance = tcli_service_class.return_value
                 results_mock = Mock()
                 results_mock.startRowOffset = 0
@@ -1003,7 +1019,7 @@ def test_handle_execute_response_reads_has_more_rows_in_direct_results(
                     resultSetMetadata=self.metadata_resp,
                     resultSet=ttypes.TFetchResultsResp(
                         status=self.okay_status,
-                        hasMoreRows=has_more_rows,
+                        hasMoreRows=is_direct_results,
                         results=results_mock,
                     ),
                     closeOperation=Mock(),
@@ -1019,11 +1035,12 @@ def test_handle_execute_response_reads_has_more_rows_in_direct_results(
                 )
                 thrift_backend = self._make_fake_thrift_backend()
 
-                execute_response, _ = thrift_backend._handle_execute_response(
-                    execute_resp, Mock()
-                )
+                (
+                    execute_response,
+                    has_more_rows_result,
+                ) = thrift_backend._handle_execute_response(execute_resp, Mock())
 
-                self.assertEqual(has_more_rows, execute_response.has_more_rows)
+                self.assertEqual(is_direct_results, has_more_rows_result)
 
     @patch(
         "databricks.sql.utils.ResultSetQueueFactory.build_queue", return_value=Mock()
@@ -1032,10 +1049,10 @@ def test_handle_execute_response_reads_has_more_rows_in_direct_results(
     def test_handle_execute_response_reads_has_more_rows_in_result_response(
         self, tcli_service_class, build_queue
     ):
-        for has_more_rows, resp_type in itertools.product(
+        for is_direct_results, resp_type in itertools.product(
             [True, False], self.execute_response_types
         ):
-            with self.subTest(has_more_rows=has_more_rows, resp_type=resp_type):
+            with self.subTest(is_direct_results=is_direct_results, resp_type=resp_type):
                 tcli_service_instance = tcli_service_class.return_value
                 results_mock = MagicMock()
                 results_mock.startRowOffset = 0
@@ -1048,7 +1065,7 @@ def test_handle_execute_response_reads_has_more_rows_in_result_response(
 
                 fetch_results_resp = ttypes.TFetchResultsResp(
                     status=self.okay_status,
-                    hasMoreRows=has_more_rows,
+                    hasMoreRows=is_direct_results,
                     results=results_mock,
                     resultSetMetadata=ttypes.TGetResultSetMetadataResp(
                         resultFormat=ttypes.TSparkRowSetType.ARROW_BASED_SET
@@ -1081,7 +1098,7 @@ def test_handle_execute_response_reads_has_more_rows_in_result_response(
                     description=Mock(),
                 )
 
-                self.assertEqual(has_more_rows, has_more_rows_resp)
+                self.assertEqual(is_direct_results, has_more_rows_resp)
 
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_arrow_batches_row_count_are_respected(self, tcli_service_class):
@@ -1136,9 +1153,10 @@ def test_arrow_batches_row_count_are_respected(self, tcli_service_class):
 
         self.assertEqual(arrow_queue.n_valid_rows, 15 * 10)
 
+    @patch("databricks.sql.backend.thrift_backend.ThriftResultSet")
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_execute_statement_calls_client_and_handle_execute_response(
-        self, tcli_service_class
+        self, tcli_service_class, mock_result_set
     ):
         tcli_service_instance = tcli_service_class.return_value
         response = Mock()
@@ -1151,14 +1169,15 @@ def test_execute_statement_calls_client_and_handle_execute_response(
             auth_provider=AuthProvider(),
             ssl_options=SSLOptions(),
         )
-        thrift_backend._handle_execute_response = Mock(return_value=(Mock(), Mock()))
+        thrift_backend._handle_execute_response = Mock()
+        thrift_backend._handle_execute_response.return_value = (Mock(), Mock())
         cursor_mock = Mock()
 
         result = thrift_backend.execute_command(
             "foo", Mock(), 100, 200, Mock(), cursor_mock
         )
         # Verify the result is a ResultSet
-        self.assertIsInstance(result, ResultSet)
+        self.assertEqual(result, mock_result_set.return_value)
 
         # Check call to client
         req = tcli_service_instance.ExecuteStatement.call_args[0][0]
@@ -1170,9 +1189,10 @@ def test_execute_statement_calls_client_and_handle_execute_response(
             response, cursor_mock
         )
 
+    @patch("databricks.sql.backend.thrift_backend.ThriftResultSet")
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_get_catalogs_calls_client_and_handle_execute_response(
-        self, tcli_service_class
+        self, tcli_service_class, mock_result_set
     ):
         tcli_service_instance = tcli_service_class.return_value
         response = Mock()
@@ -1185,12 +1205,13 @@ def test_get_catalogs_calls_client_and_handle_execute_response(
             auth_provider=AuthProvider(),
             ssl_options=SSLOptions(),
         )
-        thrift_backend._handle_execute_response = Mock(return_value=(Mock(), Mock()))
+        thrift_backend._handle_execute_response = Mock()
+        thrift_backend._handle_execute_response.return_value = (Mock(), Mock())
         cursor_mock = Mock()
 
         result = thrift_backend.get_catalogs(Mock(), 100, 200, cursor_mock)
         # Verify the result is a ResultSet
-        self.assertIsInstance(result, ResultSet)
+        self.assertEqual(result, mock_result_set.return_value)
 
         # Check call to client
         req = tcli_service_instance.GetCatalogs.call_args[0][0]
@@ -1201,9 +1222,10 @@ def test_get_catalogs_calls_client_and_handle_execute_response(
             response, cursor_mock
         )
 
+    @patch("databricks.sql.backend.thrift_backend.ThriftResultSet")
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_get_schemas_calls_client_and_handle_execute_response(
-        self, tcli_service_class
+        self, tcli_service_class, mock_result_set
     ):
         tcli_service_instance = tcli_service_class.return_value
         response = Mock()
@@ -1216,7 +1238,8 @@ def test_get_schemas_calls_client_and_handle_execute_response(
             auth_provider=AuthProvider(),
             ssl_options=SSLOptions(),
         )
-        thrift_backend._handle_execute_response = Mock(return_value=(Mock(), Mock()))
+        thrift_backend._handle_execute_response = Mock()
+        thrift_backend._handle_execute_response.return_value = (Mock(), Mock())
         cursor_mock = Mock()
 
         result = thrift_backend.get_schemas(
@@ -1228,7 +1251,7 @@ def test_get_schemas_calls_client_and_handle_execute_response(
             schema_name="schema_pattern",
         )
         # Verify the result is a ResultSet
-        self.assertIsInstance(result, ResultSet)
+        self.assertEqual(result, mock_result_set.return_value)
 
         # Check call to client
         req = tcli_service_instance.GetSchemas.call_args[0][0]
@@ -1241,9 +1264,10 @@ def test_get_schemas_calls_client_and_handle_execute_response(
             response, cursor_mock
         )
 
+    @patch("databricks.sql.backend.thrift_backend.ThriftResultSet")
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_get_tables_calls_client_and_handle_execute_response(
-        self, tcli_service_class
+        self, tcli_service_class, mock_result_set
     ):
         tcli_service_instance = tcli_service_class.return_value
         response = Mock()
@@ -1256,7 +1280,8 @@ def test_get_tables_calls_client_and_handle_execute_response(
             auth_provider=AuthProvider(),
             ssl_options=SSLOptions(),
         )
-        thrift_backend._handle_execute_response = Mock(return_value=(Mock(), Mock()))
+        thrift_backend._handle_execute_response = Mock()
+        thrift_backend._handle_execute_response.return_value = (Mock(), Mock())
         cursor_mock = Mock()
 
         result = thrift_backend.get_tables(
@@ -1270,7 +1295,7 @@ def test_get_tables_calls_client_and_handle_execute_response(
             table_types=["type1", "type2"],
         )
         # Verify the result is a ResultSet
-        self.assertIsInstance(result, ResultSet)
+        self.assertEqual(result, mock_result_set.return_value)
 
         # Check call to client
         req = tcli_service_instance.GetTables.call_args[0][0]
@@ -1285,9 +1310,10 @@ def test_get_tables_calls_client_and_handle_execute_response(
             response, cursor_mock
         )
 
+    @patch("databricks.sql.backend.thrift_backend.ThriftResultSet")
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_get_columns_calls_client_and_handle_execute_response(
-        self, tcli_service_class
+        self, tcli_service_class, mock_result_set
     ):
         tcli_service_instance = tcli_service_class.return_value
         response = Mock()
@@ -1300,7 +1326,8 @@ def test_get_columns_calls_client_and_handle_execute_response(
             auth_provider=AuthProvider(),
             ssl_options=SSLOptions(),
         )
-        thrift_backend._handle_execute_response = Mock(return_value=(Mock(), Mock()))
+        thrift_backend._handle_execute_response = Mock()
+        thrift_backend._handle_execute_response.return_value = (Mock(), Mock())
         cursor_mock = Mock()
 
         result = thrift_backend.get_columns(
@@ -1314,7 +1341,7 @@ def test_get_columns_calls_client_and_handle_execute_response(
             column_name="column_pattern",
         )
         # Verify the result is a ResultSet
-        self.assertIsInstance(result, ResultSet)
+        self.assertEqual(result, mock_result_set.return_value)
 
         # Check call to client
         req = tcli_service_instance.GetColumns.call_args[0][0]
@@ -2203,14 +2230,23 @@ def test_protocol_v3_fails_if_initial_namespace_set(self, tcli_client_class):
             str(cm.exception),
         )
 
+    @patch("databricks.sql.backend.thrift_backend.ThriftResultSet")
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     @patch(
         "databricks.sql.backend.thrift_backend.ThriftDatabricksClient._handle_execute_response"
     )
     def test_execute_command_sets_complex_type_fields_correctly(
-        self, mock_handle_execute_response, tcli_service_class
+        self, mock_handle_execute_response, tcli_service_class, mock_result_set
     ):
         tcli_service_instance = tcli_service_class.return_value
+        # Set up the mock to return a tuple with two values
+        mock_execute_response = Mock()
+        mock_arrow_schema = Mock()
+        mock_handle_execute_response.return_value = (
+            mock_execute_response,
+            mock_arrow_schema,
+        )
+
         # Iterate through each possible combination of native types (True, False and unset)
         for complex, timestamp, decimals in itertools.product(
             [True, False, None], [True, False, None], [True, False, None]