@@ -619,13 +619,6 @@ def test_handle_execute_response_sets_compression_in_direct_results(
619
619
lz4Compressed = Mock ()
620
620
resultSet = MagicMock ()
621
621
resultSet .results .startRowOffset = 0
622
-
623
- # Create a valid operation status
624
- op_status = ttypes .TGetOperationStatusResp (
625
- status = self .okay_status ,
626
- operationState = ttypes .TOperationState .FINISHED_STATE ,
627
- )
628
-
629
622
t_execute_resp = resp_type (
630
623
status = Mock (),
631
624
operationHandle = Mock (),
@@ -927,9 +920,7 @@ def test_handle_execute_response_can_handle_with_direct_results(self):
927
920
auth_provider = AuthProvider (),
928
921
ssl_options = SSLOptions (),
929
922
)
930
- thrift_backend ._results_message_to_execute_response = Mock (
931
- return_value = (Mock (), Mock ())
932
- )
923
+ thrift_backend ._results_message_to_execute_response = Mock ()
933
924
934
925
thrift_backend ._handle_execute_response (execute_resp , Mock ())
935
926
@@ -957,12 +948,6 @@ def test_use_arrow_schema_if_available(self, tcli_service_class):
957
948
operationHandle = self .operation_handle ,
958
949
)
959
950
960
- # Mock the operation status response
961
- op_state = ttypes .TGetOperationStatusResp (
962
- status = self .okay_status ,
963
- operationState = ttypes .TOperationState .FINISHED_STATE ,
964
- )
965
- tcli_service_instance .GetOperationStatus .return_value = op_state
966
951
tcli_service_instance .GetResultSetMetadata .return_value = (
967
952
t_get_result_set_metadata_resp
968
953
)
@@ -977,7 +962,7 @@ def test_use_arrow_schema_if_available(self, tcli_service_class):
977
962
t_execute_resp , Mock ()
978
963
)
979
964
980
- self .assertEqual (arrow_schema_bytes , arrow_schema_mock )
965
+ self .assertEqual (execute_response . arrow_schema_bytes , arrow_schema_mock )
981
966
982
967
@patch ("databricks.sql.backend.thrift_backend.TCLIService.Client" , autospec = True )
983
968
def test_fall_back_to_hive_schema_if_no_arrow_schema (self , tcli_service_class ):
@@ -997,12 +982,6 @@ def test_fall_back_to_hive_schema_if_no_arrow_schema(self, tcli_service_class):
997
982
operationHandle = self .operation_handle ,
998
983
)
999
984
1000
- # Mock the operation status response
1001
- op_state = ttypes .TGetOperationStatusResp (
1002
- status = self .okay_status ,
1003
- operationState = ttypes .TOperationState .FINISHED_STATE ,
1004
- )
1005
- tcli_service_instance .GetOperationStatus .return_value = op_state
1006
985
tcli_service_instance .GetResultSetMetadata .return_value = hive_schema_req
1007
986
tcli_service_instance .GetOperationStatus .return_value = (
1008
987
ttypes .TGetOperationStatusResp (
@@ -1694,9 +1673,7 @@ def test_handle_execute_response_sets_active_op_handle(self):
1694
1673
thrift_backend = self ._make_fake_thrift_backend ()
1695
1674
thrift_backend ._check_direct_results_for_error = Mock ()
1696
1675
thrift_backend ._wait_until_command_done = Mock ()
1697
- thrift_backend ._results_message_to_execute_response = Mock (
1698
- return_value = (Mock (), Mock ())
1699
- )
1676
+ thrift_backend ._results_message_to_execute_response = Mock ()
1700
1677
1701
1678
# Create a mock response with a real operation handle
1702
1679
mock_resp = Mock ()
@@ -2256,8 +2233,7 @@ def test_protocol_v3_fails_if_initial_namespace_set(self, tcli_client_class):
2256
2233
@patch ("databricks.sql.backend.thrift_backend.ThriftResultSet" )
2257
2234
@patch ("databricks.sql.backend.thrift_backend.TCLIService.Client" , autospec = True )
2258
2235
@patch (
2259
- "databricks.sql.backend.thrift_backend.ThriftDatabricksClient._handle_execute_response" ,
2260
- return_value = (Mock (), Mock ()),
2236
+ "databricks.sql.backend.thrift_backend.ThriftDatabricksClient._handle_execute_response"
2261
2237
)
2262
2238
def test_execute_command_sets_complex_type_fields_correctly (
2263
2239
self , mock_handle_execute_response , tcli_service_class , mock_result_set
0 commit comments