Configure autospec for mocked Client objects (#188)

Resolves #187

Signed-off-by: Jesse Whitehouse <[email protected]>

Author: Jesse

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## 2.8.x (Unreleased)
 
 - Other: Fix typo in README quick start example
+- Other: Add autospec to Client mocks and tidy up `make_request`
 
 ## 2.8.0 (2023-07-21)
 

src/databricks/sql/thrift_backend.py

@@ -335,12 +335,9 @@ def attempt_request(attempt):
 
             error, error_message, retry_delay = None, None, None
             try:
-                # The MagicMocks in our unit tests have a `name` property instead of `__name__`.
                 logger.debug(
                     "Sending request: {}(<REDACTED>)".format(
-                        getattr(
-                            method, "__name__", getattr(method, "name", "UnknownMethod")
-                        )
+                        getattr(method, "__name__")
                     )
                 )
                 unsafe_logger.debug("Sending request: {}".format(request))

tests/unit/test_thrift_backend.py

@@ -96,7 +96,7 @@ def test_hive_schema_to_arrow_schema_preserves_column_names(self):
         self.assertEqual(arrow_schema.field(2).name, "column 2")
         self.assertEqual(arrow_schema.field(3).name, "")
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_bad_protocol_versions_are_rejected(self, tcli_service_client_cass):
         t_http_client_instance = tcli_service_client_cass.return_value
         bad_protocol_versions = [
@@ -123,7 +123,7 @@ def test_bad_protocol_versions_are_rejected(self, tcli_service_client_cass):
 
             self.assertIn("expected server to use a protocol version", str(cm.exception))
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_okay_protocol_versions_succeed(self, tcli_service_client_cass):
         t_http_client_instance = tcli_service_client_cass.return_value
         good_protocol_versions = [
@@ -351,7 +351,7 @@ def test_handle_execute_response_sets_compression_in_direct_results(self, build_
             execute_response = thrift_backend._handle_execute_response(t_execute_resp, Mock())
             self.assertEqual(execute_response.lz4_compressed, lz4Compressed)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_handle_execute_response_checks_operation_state_in_polls(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
 
@@ -380,7 +380,7 @@ def test_handle_execute_response_checks_operation_state_in_polls(self, tcli_serv
                 if op_state_resp.errorMessage:
                     self.assertIn(op_state_resp.errorMessage, str(cm.exception))
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_get_status_uses_display_message_if_available(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
 
@@ -405,7 +405,7 @@ def test_get_status_uses_display_message_if_available(self, tcli_service_class):
         self.assertEqual(display_message, str(cm.exception))
         self.assertIn(diagnostic_info, str(cm.exception.message_with_context()))
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_direct_results_uses_display_message_if_available(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
 
@@ -477,7 +477,7 @@ 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.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_handle_execute_response_can_handle_without_direct_results(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
 
@@ -542,7 +542,7 @@ def test_handle_execute_response_can_handle_with_direct_results(self):
                     ttypes.TOperationState.FINISHED_STATE,
                 )
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_use_arrow_schema_if_available(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
         arrow_schema_mock = MagicMock(name="Arrow schema mock")
@@ -566,7 +566,7 @@ def test_use_arrow_schema_if_available(self, tcli_service_class):
 
         self.assertEqual(execute_response.arrow_schema_bytes, arrow_schema_mock)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_fall_back_to_hive_schema_if_no_arrow_schema(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
         hive_schema_mock = MagicMock(name="Hive schema mock")
@@ -591,7 +591,7 @@ def test_fall_back_to_hive_schema_if_no_arrow_schema(self, tcli_service_class):
                          thrift_backend._hive_schema_to_arrow_schema.call_args[0][0])
 
     @patch("databricks.sql.utils.ResultSetQueueFactory.build_queue", return_value=Mock())
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     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([True, False],
@@ -625,7 +625,7 @@ def test_handle_execute_response_reads_has_more_rows_in_direct_results(
                 self.assertEqual(has_more_rows, execute_response.has_more_rows)
 
     @patch("databricks.sql.utils.ResultSetQueueFactory.build_queue", return_value=Mock())
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     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([True, False],
@@ -671,7 +671,7 @@ def test_handle_execute_response_reads_has_more_rows_in_result_response(
 
                 self.assertEqual(has_more_rows, has_more_rows_resp)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_arrow_batches_row_count_are_respected(self, tcli_service_class):
         # make some semi-real arrow batches and check the number of rows is correct in the queue
         tcli_service_instance = tcli_service_class.return_value
@@ -709,7 +709,7 @@ def test_arrow_batches_row_count_are_respected(self, tcli_service_class):
 
         self.assertEqual(arrow_queue.n_valid_rows, 15 * 10)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_execute_statement_calls_client_and_handle_execute_response(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
         response = Mock()
@@ -727,7 +727,7 @@ def test_execute_statement_calls_client_and_handle_execute_response(self, tcli_s
         # Check response handling
         thrift_backend._handle_execute_response.assert_called_with(response, cursor_mock)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_get_catalogs_calls_client_and_handle_execute_response(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
         response = Mock()
@@ -744,7 +744,7 @@ def test_get_catalogs_calls_client_and_handle_execute_response(self, tcli_servic
         # Check response handling
         thrift_backend._handle_execute_response.assert_called_with(response, cursor_mock)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_get_schemas_calls_client_and_handle_execute_response(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
         response = Mock()
@@ -769,7 +769,7 @@ def test_get_schemas_calls_client_and_handle_execute_response(self, tcli_service
         # Check response handling
         thrift_backend._handle_execute_response.assert_called_with(response, cursor_mock)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_get_tables_calls_client_and_handle_execute_response(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
         response = Mock()
@@ -798,7 +798,7 @@ def test_get_tables_calls_client_and_handle_execute_response(self, tcli_service_
         # Check response handling
         thrift_backend._handle_execute_response.assert_called_with(response, cursor_mock)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_get_columns_calls_client_and_handle_execute_response(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
         response = Mock()
@@ -827,7 +827,7 @@ def test_get_columns_calls_client_and_handle_execute_response(self, tcli_service
         # Check response handling
         thrift_backend._handle_execute_response.assert_called_with(response, cursor_mock)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_open_session_user_provided_session_id_optional(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
         tcli_service_instance.OpenSession.return_value = self.open_session_resp
@@ -836,23 +836,23 @@ def test_open_session_user_provided_session_id_optional(self, tcli_service_class
         thrift_backend.open_session({}, None, None)
         self.assertEqual(len(tcli_service_instance.OpenSession.call_args_list), 1)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_op_handle_respected_in_close_command(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
         thrift_backend = ThriftBackend("foobar", 443, "path", [], auth_provider=AuthProvider())
         thrift_backend.close_command(self.operation_handle)
         self.assertEqual(tcli_service_instance.CloseOperation.call_args[0][0].operationHandle,
                          self.operation_handle)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_session_handle_respected_in_close_session(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
         thrift_backend = ThriftBackend("foobar", 443, "path", [], auth_provider=AuthProvider())
         thrift_backend.close_session(self.session_handle)
         self.assertEqual(tcli_service_instance.CloseSession.call_args[0][0].sessionHandle,
                          self.session_handle)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_non_arrow_non_column_based_set_triggers_exception(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
         results_mock = Mock()
@@ -1021,7 +1021,7 @@ def test_convert_column_based_set_to_arrow_table_uses_types_from_col_set(self):
         self.assertEqual(arrow_table.column(2).to_pylist(), [1.15, 2.2, 3.3])
         self.assertEqual(arrow_table.column(3).to_pylist(), [b'\x11', b'\x22', b'\x33'])
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_cancel_command_uses_active_op_handle(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
 
@@ -1318,7 +1318,7 @@ def test_retry_args_bounding(self, mock_http_client):
             for (arg, val) in retry_delay_expected_vals.items():
                 self.assertEqual(getattr(backend, arg), val)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_configuration_passthrough(self, tcli_client_class):
         tcli_service_instance = tcli_client_class.return_value
         tcli_service_instance.OpenSession.return_value = self.open_session_resp
@@ -1336,7 +1336,7 @@ def test_configuration_passthrough(self, tcli_client_class):
         open_session_req = tcli_client_class.return_value.OpenSession.call_args[0][0]
         self.assertEqual(open_session_req.configuration, expected_config)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_cant_set_timestamp_as_string_to_true(self, tcli_client_class):
         tcli_service_instance = tcli_client_class.return_value
         tcli_service_instance.OpenSession.return_value = self.open_session_resp
@@ -1355,7 +1355,7 @@ def _construct_open_session_with_namespace(self, can_use_multiple_cats, cat, sch
             canUseMultipleCatalogs=can_use_multiple_cats,
             initialNamespace=ttypes.TNamespace(catalogName=cat, schemaName=schem))
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_initial_namespace_passthrough_to_open_session(self, tcli_client_class):
         tcli_service_instance = tcli_client_class.return_value
 
@@ -1373,7 +1373,7 @@ def test_initial_namespace_passthrough_to_open_session(self, tcli_client_class):
                 self.assertEqual(open_session_req.initialNamespace.catalogName, cat)
                 self.assertEqual(open_session_req.initialNamespace.schemaName, schem)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_can_use_multiple_catalogs_is_set_in_open_session_req(self, tcli_client_class):
         tcli_service_instance = tcli_client_class.return_value
         tcli_service_instance.OpenSession.return_value = self.open_session_resp
@@ -1384,7 +1384,7 @@ def test_can_use_multiple_catalogs_is_set_in_open_session_req(self, tcli_client_
         open_session_req = tcli_client_class.return_value.OpenSession.call_args[0][0]
         self.assertTrue(open_session_req.canUseMultipleCatalogs)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_can_use_multiple_catalogs_is_false_fails_with_initial_catalog(self, tcli_client_class):
         tcli_service_instance = tcli_client_class.return_value
 
@@ -1410,7 +1410,7 @@ def test_can_use_multiple_catalogs_is_false_fails_with_initial_catalog(self, tcl
                 self._construct_open_session_with_namespace(False, cat, schem)
             backend.open_session({}, cat, schem)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     def test_protocol_v3_fails_if_initial_namespace_set(self, tcli_client_class):
         tcli_service_instance = tcli_client_class.return_value
 
@@ -1430,7 +1430,7 @@ def test_protocol_v3_fails_if_initial_namespace_set(self, tcli_client_class):
         self.assertIn("Setting initial namespace not supported by the DBR version",
                       str(cm.exception))
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
     @patch("databricks.sql.thrift_backend.ThriftBackend._handle_execute_response")
     def test_execute_command_sets_complex_type_fields_correctly(self, mock_handle_execute_response,
                                                                 tcli_service_class):