Skip to content

Commit 78cfd27

Browse files
formatting (black)
1 parent 9a968ed commit 78cfd27

File tree

2 files changed

+43
-21
lines changed

2 files changed

+43
-21
lines changed

tests/unit/test_client.py

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ def new(cls):
3535
ThriftBackendMock.return_value = ThriftBackendMock
3636

3737
cls.apply_property_to_mock(ThriftBackendMock, staging_allowed_local_path=None)
38-
38+
3939
mock_result_set = Mock(spec=ThriftResultSet)
4040
cls.apply_property_to_mock(
4141
mock_result_set,
@@ -80,21 +80,24 @@ class ClientTestSuite(unittest.TestCase):
8080
"access_token": "tok",
8181
}
8282

83-
@patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME, ThriftDatabricksClientMockFactory.new())
83+
@patch(
84+
"%s.session.ThriftDatabricksClient" % PACKAGE_NAME,
85+
ThriftDatabricksClientMockFactory.new(),
86+
)
8487
def test_closing_connection_closes_commands(self):
8588
for closed in (True, False):
8689
with self.subTest(closed=closed):
8790
connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
8891
cursor = connection.cursor()
89-
92+
9093
# Create a mock result set and set it as the active result set
9194
mock_result_set = Mock()
9295
mock_result_set.has_been_closed_server_side = closed
9396
cursor.active_result_set = mock_result_set
94-
97+
9598
# Close the connection
9699
connection.close()
97-
100+
98101
# After closing the connection, the close method should have been called on the result set
99102
mock_result_set.close.assert_called_once_with()
100103

@@ -167,15 +170,13 @@ def test_executing_multiple_commands_uses_the_most_recent_command(
167170
# Set is_staging_operation to False to avoid _handle_staging_operation being called
168171
for mock_rs in mock_result_sets:
169172
mock_rs.is_staging_operation = False
170-
173+
171174
mock_result_set_class.side_effect = mock_result_sets
172-
175+
173176
mock_backend = ThriftDatabricksClientMockFactory.new()
174177
mock_backend.execute_command.side_effect = mock_result_sets
175178

176-
cursor = client.Cursor(
177-
connection=Mock(), backend=mock_backend
178-
)
179+
cursor = client.Cursor(connection=Mock(), backend=mock_backend)
179180
cursor.execute("SELECT 1;")
180181
cursor.execute("SELECT 1;")
181182

@@ -351,7 +352,7 @@ def test_executemany_parameter_passhthrough_and_uses_last_result_set(
351352
# Set is_staging_operation to False to avoid _handle_staging_operation being called
352353
for mock_rs in mock_result_set_instances:
353354
mock_rs.is_staging_operation = False
354-
355+
355356
mock_result_set_class.side_effect = mock_result_set_instances
356357
mock_backend = ThriftDatabricksClientMockFactory.new()
357358
mock_backend.execute_command.side_effect = mock_result_set_instances
@@ -518,7 +519,10 @@ def test_staging_operation_response_is_handled(
518519

519520
mock_handle_staging_operation.call_count == 1
520521

521-
@patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME, ThriftDatabricksClientMockFactory.new())
522+
@patch(
523+
"%s.session.ThriftDatabricksClient" % PACKAGE_NAME,
524+
ThriftDatabricksClientMockFactory.new(),
525+
)
522526
def test_access_current_query_id(self):
523527
operation_id = "EE6A8778-21FC-438B-92D8-96AC51EE3821"
524528

tests/unit/test_thrift_backend.py

Lines changed: 27 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,9 @@ def test_hive_schema_to_arrow_schema_preserves_column_names(self):
127127
]
128128

129129
t_table_schema = ttypes.TTableSchema(columns)
130-
arrow_schema = ThriftDatabricksClient._hive_schema_to_arrow_schema(t_table_schema)
130+
arrow_schema = ThriftDatabricksClient._hive_schema_to_arrow_schema(
131+
t_table_schema
132+
)
131133

132134
self.assertEqual(arrow_schema.field(0).name, "column 1")
133135
self.assertEqual(arrow_schema.field(1).name, "column 2")
@@ -1147,7 +1149,9 @@ def test_execute_statement_calls_client_and_handle_execute_response(
11471149
thrift_backend._handle_execute_response = Mock()
11481150
cursor_mock = Mock()
11491151

1150-
result = thrift_backend.execute_command("foo", Mock(), 100, 200, Mock(), cursor_mock)
1152+
result = thrift_backend.execute_command(
1153+
"foo", Mock(), 100, 200, Mock(), cursor_mock
1154+
)
11511155
# Verify the result is a ResultSet
11521156
self.assertIsInstance(result, ResultSet)
11531157

@@ -1419,8 +1423,12 @@ def test_create_arrow_table_raises_error_for_unsupported_type(self):
14191423
with self.assertRaises(OperationalError):
14201424
thrift_backend._create_arrow_table(t_row_set, Mock(), None, Mock())
14211425

1422-
@patch("databricks.sql.backend.thrift_backend.convert_arrow_based_set_to_arrow_table")
1423-
@patch("databricks.sql.backend.thrift_backend.convert_column_based_set_to_arrow_table")
1426+
@patch(
1427+
"databricks.sql.backend.thrift_backend.convert_arrow_based_set_to_arrow_table"
1428+
)
1429+
@patch(
1430+
"databricks.sql.backend.thrift_backend.convert_column_based_set_to_arrow_table"
1431+
)
14241432
def test_create_arrow_table_calls_correct_conversion_method(
14251433
self, convert_col_mock, convert_arrow_mock
14261434
):
@@ -1643,7 +1651,8 @@ def test_handle_execute_response_sets_active_op_handle(self):
16431651
"databricks.sql.thrift_api.TCLIService.TCLIService.Client.GetOperationStatus"
16441652
)
16451653
@patch(
1646-
"databricks.sql.backend.thrift_backend._retry_policy", new_callable=retry_policy_factory
1654+
"databricks.sql.backend.thrift_backend._retry_policy",
1655+
new_callable=retry_policy_factory,
16471656
)
16481657
def test_make_request_will_retry_GetOperationStatus(
16491658
self, mock_retry_policy, mock_GetOperationStatus, t_transport_class
@@ -1718,7 +1727,8 @@ def test_make_request_will_retry_GetOperationStatus(
17181727
"databricks.sql.thrift_api.TCLIService.TCLIService.Client.GetOperationStatus"
17191728
)
17201729
@patch(
1721-
"databricks.sql.backend.thrift_backend._retry_policy", new_callable=retry_policy_factory
1730+
"databricks.sql.backend.thrift_backend._retry_policy",
1731+
new_callable=retry_policy_factory,
17221732
)
17231733
def test_make_request_will_retry_GetOperationStatus_for_http_error(
17241734
self, mock_retry_policy, mock_gos
@@ -1795,7 +1805,8 @@ def test_make_request_wont_retry_if_error_code_not_429_or_503(
17951805

17961806
@patch("databricks.sql.auth.thrift_http_client.THttpClient")
17971807
@patch(
1798-
"databricks.sql.backend.thrift_backend._retry_policy", new_callable=retry_policy_factory
1808+
"databricks.sql.backend.thrift_backend._retry_policy",
1809+
new_callable=retry_policy_factory,
17991810
)
18001811
def test_make_request_will_retry_stop_after_attempts_count_if_retryable(
18011812
self, mock_retry_policy, t_transport_class
@@ -1975,7 +1986,12 @@ def test_retry_args_passthrough(self, mock_http_client):
19751986
@patch("thrift.transport.THttpClient.THttpClient")
19761987
def test_retry_args_bounding(self, mock_http_client):
19771988
retry_delay_test_args_and_expected_values = {}
1978-
for k, (_, _, min, max) in databricks.sql.backend.thrift_backend._retry_policy.items():
1989+
for k, (
1990+
_,
1991+
_,
1992+
min,
1993+
max,
1994+
) in databricks.sql.backend.thrift_backend._retry_policy.items():
19791995
retry_delay_test_args_and_expected_values[k] = (
19801996
(min - 1, min),
19811997
(max + 1, max),
@@ -2171,7 +2187,9 @@ def test_protocol_v3_fails_if_initial_namespace_set(self, tcli_client_class):
21712187
)
21722188

21732189
@patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
2174-
@patch("databricks.sql.backend.thrift_backend.ThriftDatabricksClient._handle_execute_response")
2190+
@patch(
2191+
"databricks.sql.backend.thrift_backend.ThriftDatabricksClient._handle_execute_response"
2192+
)
21752193
def test_execute_command_sets_complex_type_fields_correctly(
21762194
self, mock_handle_execute_response, tcli_service_class
21772195
):

0 commit comments

Comments
 (0)