Skip to content

Configure autospec for mocked Client objects #188

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
susodapop merged 2 commits into from
Aug 9, 2023
Merged

Configure autospec for mocked Client objects #188

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
894c748
Configure autospec for mocked Client objects
Aug 9, 2023
ea7ae74
Update changelog
Aug 9, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
5 changes: 1 addition & 4 deletions src/databricks/sql/thrift_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
56 changes: 28 additions & 28 deletions tests/unit/test_thrift_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand All @@ -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 = [
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand All @@ -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

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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")
Expand All @@ -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")
Expand All @@ -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],
Expand Down Expand Up @@ -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],
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand All @@ -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()
Expand All @@ -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()
Expand All @@ -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()
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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
Expand All @@ -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()
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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

Expand All @@ -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
Expand All @@ -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

Expand All @@ -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

Expand All @@ -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):
Expand Down