call correct method in test_use_arrow_schema

varun-edachali-dbx · varun-edachali-dbx · commit 5ee413677016 · 2025-06-10T13:14:18.000Z
Signed-off-by: varun-edachali-dbx &lt;varun.edachali@databricks.com&gt;
diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py
@@ -966,8 +966,8 @@ def test_use_arrow_schema_if_available(self, tcli_service_class):
         thrift_backend = self._create_thrift_client()
 
         # Call the real _results_message_to_execute_response method
-        execute_response, _ = thrift_backend._results_message_to_execute_response(
-            t_execute_resp, ttypes.TOperationState.FINISHED_STATE
+        execute_response, _ = thrift_backend._handle_execute_response(
+            t_execute_resp, Mock()
         )
 
         self.assertEqual(execute_response.arrow_schema_bytes, arrow_schema_mock)

Original file line number	Diff line number	Diff line change
`@@ -966,8 +966,8 @@ def test_use_arrow_schema_if_available(self, tcli_service_class):`
`966`	`966`	`thrift_backend = self._create_thrift_client()`
`967`	`967`
`968`	`968`	`# Call the real _results_message_to_execute_response method`
`969`		`- execute_response, _ = thrift_backend._results_message_to_execute_response(`
`970`		`- t_execute_resp, ttypes.TOperationState.FINISHED_STATE`
	`969`	`+ execute_response, _ = thrift_backend._handle_execute_response(`
	`970`	`+ t_execute_resp, Mock()`
`971`	`971`	`)`
`972`	`972`
`973`	`973`	`self.assertEqual(execute_response.arrow_schema_bytes, arrow_schema_mock)`