@@ -1009,13 +1009,12 @@ def test_fall_back_to_hive_schema_if_no_arrow_schema(self, tcli_service_class):
10091009 "databricks.sql.utils.ResultSetQueueFactory.build_queue" , return_value = Mock ()
10101010 )
10111011 @patch ("databricks.sql.backend.thrift_backend.TCLIService.Client" , autospec = True )
1012- def test_handle_execute_response_reads_has_more_rows_in_direct_results (
1012+ def test_handle_execute_response_creates_execute_response (
10131013 self , tcli_service_class , build_queue
10141014 ):
1015- for has_more_rows , resp_type in itertools .product (
1016- [True , False ], self .execute_response_types
1017- ):
1018- with self .subTest (has_more_rows = has_more_rows , resp_type = resp_type ):
1015+ """Test that _handle_execute_response creates an ExecuteResponse object correctly."""
1016+ for resp_type in self .execute_response_types :
1017+ with self .subTest (resp_type = resp_type ):
10191018 tcli_service_instance = tcli_service_class .return_value
10201019 results_mock = Mock ()
10211020 results_mock .startRowOffset = 0
@@ -1027,7 +1026,7 @@ def test_handle_execute_response_reads_has_more_rows_in_direct_results(
10271026 resultSetMetadata = self .metadata_resp ,
10281027 resultSet = ttypes .TFetchResultsResp (
10291028 status = self .okay_status ,
1030- hasMoreRows = has_more_rows ,
1029+ hasMoreRows = True ,
10311030 results = results_mock ,
10321031 ),
10331032 closeOperation = Mock (),
@@ -1047,7 +1046,8 @@ def test_handle_execute_response_reads_has_more_rows_in_direct_results(
10471046 execute_resp , Mock ()
10481047 )
10491048
1050- self .assertEqual (has_more_rows , execute_response .has_more_rows )
1049+ self .assertIsNotNone (execute_response )
1050+ self .assertIsInstance (execute_response , ExecuteResponse )
10511051
10521052 @patch (
10531053 "databricks.sql.utils.ResultSetQueueFactory.build_queue" , return_value = Mock ()
0 commit comments