@@ -990,7 +990,6 @@ def test_fall_back_to_hive_schema_if_no_arrow_schema(self, tcli_service_class):
990990 operationHandle = self .operation_handle ,
991991 )
992992
993- # Mock the operation status response
994993 op_state = ttypes .TGetOperationStatusResp (
995994 status = self .okay_status ,
996995 operationState = ttypes .TOperationState .FINISHED_STATE ,
@@ -1011,52 +1010,6 @@ def test_fall_back_to_hive_schema_if_no_arrow_schema(self, tcli_service_class):
10111010 thrift_backend ._hive_schema_to_arrow_schema .call_args [0 ][0 ],
10121011 )
10131012
1014- @patch (
1015- "databricks.sql.utils.ResultSetQueueFactory.build_queue" , return_value = Mock ()
1016- )
1017- @patch ("databricks.sql.backend.thrift_backend.TCLIService.Client" , autospec = True )
1018- def test_handle_execute_response_creates_execute_response (
1019- self , tcli_service_class , build_queue
1020- ):
1021- """Test that _handle_execute_response creates an ExecuteResponse object correctly."""
1022- for resp_type in self .execute_response_types :
1023- with self .subTest (resp_type = resp_type ):
1024- tcli_service_instance = tcli_service_class .return_value
1025- results_mock = Mock ()
1026- results_mock .startRowOffset = 0
1027- direct_results_message = ttypes .TSparkDirectResults (
1028- operationStatus = ttypes .TGetOperationStatusResp (
1029- status = self .okay_status ,
1030- operationState = ttypes .TOperationState .FINISHED_STATE ,
1031- ),
1032- resultSetMetadata = self .metadata_resp ,
1033- resultSet = ttypes .TFetchResultsResp (
1034- status = self .okay_status ,
1035- hasMoreRows = True ,
1036- results = results_mock ,
1037- ),
1038- closeOperation = Mock (),
1039- )
1040- execute_resp = resp_type (
1041- status = self .okay_status ,
1042- directResults = direct_results_message ,
1043- operationHandle = self .operation_handle ,
1044- )
1045-
1046- tcli_service_instance .GetResultSetMetadata .return_value = (
1047- self .metadata_resp
1048- )
1049- thrift_backend = self ._create_fake_thrift_client ()
1050-
1051- execute_response_tuple = thrift_backend ._handle_execute_response (
1052- execute_resp , Mock ()
1053- )
1054-
1055- self .assertIsNotNone (execute_response_tuple )
1056- self .assertIsInstance (execute_response_tuple , tuple )
1057- self .assertIsInstance (execute_response_tuple [0 ], ExecuteResponse )
1058- self .assertIsInstance (execute_response_tuple [1 ], bool )
1059-
10601013 @patch (
10611014 "databricks.sql.utils.ResultSetQueueFactory.build_queue" , return_value = Mock ()
10621015 )
0 commit comments