formatting (black)

Author: varun-edachali-dbx

tests/unit/test_client.py

@@ -35,7 +35,7 @@ def new(cls):
         ThriftBackendMock.return_value = ThriftBackendMock
 
         cls.apply_property_to_mock(ThriftBackendMock, staging_allowed_local_path=None)
-        
+
         mock_result_set = Mock(spec=ThriftResultSet)
         cls.apply_property_to_mock(
             mock_result_set,
@@ -80,21 +80,24 @@ class ClientTestSuite(unittest.TestCase):
         "access_token": "tok",
     }
 
-    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME, ThriftDatabricksClientMockFactory.new())
+    @patch(
+        "%s.session.ThriftDatabricksClient" % PACKAGE_NAME,
+        ThriftDatabricksClientMockFactory.new(),
+    )
     def test_closing_connection_closes_commands(self):
         for closed in (True, False):
             with self.subTest(closed=closed):
                 connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
                 cursor = connection.cursor()
-                
+
                 # Create a mock result set and set it as the active result set
                 mock_result_set = Mock()
                 mock_result_set.has_been_closed_server_side = closed
                 cursor.active_result_set = mock_result_set
-                
+
                 # Close the connection
                 connection.close()
-                
+
                 # After closing the connection, the close method should have been called on the result set
                 mock_result_set.close.assert_called_once_with()
 
@@ -167,15 +170,13 @@ def test_executing_multiple_commands_uses_the_most_recent_command(
         # Set is_staging_operation to False to avoid _handle_staging_operation being called
         for mock_rs in mock_result_sets:
             mock_rs.is_staging_operation = False
-            
+
         mock_result_set_class.side_effect = mock_result_sets
-        
+
         mock_backend = ThriftDatabricksClientMockFactory.new()
         mock_backend.execute_command.side_effect = mock_result_sets
 
-        cursor = client.Cursor(
-            connection=Mock(), backend=mock_backend
-        )
+        cursor = client.Cursor(connection=Mock(), backend=mock_backend)
         cursor.execute("SELECT 1;")
         cursor.execute("SELECT 1;")
 
@@ -351,7 +352,7 @@ def test_executemany_parameter_passhthrough_and_uses_last_result_set(
         # Set is_staging_operation to False to avoid _handle_staging_operation being called
         for mock_rs in mock_result_set_instances:
             mock_rs.is_staging_operation = False
-            
+
         mock_result_set_class.side_effect = mock_result_set_instances
         mock_backend = ThriftDatabricksClientMockFactory.new()
         mock_backend.execute_command.side_effect = mock_result_set_instances
@@ -518,7 +519,10 @@ def test_staging_operation_response_is_handled(
 
         mock_handle_staging_operation.call_count == 1
 
-    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME, ThriftDatabricksClientMockFactory.new())
+    @patch(
+        "%s.session.ThriftDatabricksClient" % PACKAGE_NAME,
+        ThriftDatabricksClientMockFactory.new(),
+    )
     def test_access_current_query_id(self):
         operation_id = "EE6A8778-21FC-438B-92D8-96AC51EE3821"
 

tests/unit/test_thrift_backend.py

@@ -127,7 +127,9 @@ def test_hive_schema_to_arrow_schema_preserves_column_names(self):
         ]
 
         t_table_schema = ttypes.TTableSchema(columns)
-        arrow_schema = ThriftDatabricksClient._hive_schema_to_arrow_schema(t_table_schema)
+        arrow_schema = ThriftDatabricksClient._hive_schema_to_arrow_schema(
+            t_table_schema
+        )
 
         self.assertEqual(arrow_schema.field(0).name, "column 1")
         self.assertEqual(arrow_schema.field(1).name, "column 2")
@@ -1147,7 +1149,9 @@ def test_execute_statement_calls_client_and_handle_execute_response(
         thrift_backend._handle_execute_response = Mock()
         cursor_mock = Mock()
 
-        result = thrift_backend.execute_command("foo", Mock(), 100, 200, Mock(), cursor_mock)
+        result = thrift_backend.execute_command(
+            "foo", Mock(), 100, 200, Mock(), cursor_mock
+        )
         # Verify the result is a ResultSet
         self.assertIsInstance(result, ResultSet)
 
@@ -1419,8 +1423,12 @@ def test_create_arrow_table_raises_error_for_unsupported_type(self):
         with self.assertRaises(OperationalError):
             thrift_backend._create_arrow_table(t_row_set, Mock(), None, Mock())
 
-    @patch("databricks.sql.backend.thrift_backend.convert_arrow_based_set_to_arrow_table")
-    @patch("databricks.sql.backend.thrift_backend.convert_column_based_set_to_arrow_table")
+    @patch(
+        "databricks.sql.backend.thrift_backend.convert_arrow_based_set_to_arrow_table"
+    )
+    @patch(
+        "databricks.sql.backend.thrift_backend.convert_column_based_set_to_arrow_table"
+    )
     def test_create_arrow_table_calls_correct_conversion_method(
         self, convert_col_mock, convert_arrow_mock
     ):
@@ -1643,7 +1651,8 @@ def test_handle_execute_response_sets_active_op_handle(self):
         "databricks.sql.thrift_api.TCLIService.TCLIService.Client.GetOperationStatus"
     )
     @patch(
-        "databricks.sql.backend.thrift_backend._retry_policy", new_callable=retry_policy_factory
+        "databricks.sql.backend.thrift_backend._retry_policy",
+        new_callable=retry_policy_factory,
     )
     def test_make_request_will_retry_GetOperationStatus(
         self, mock_retry_policy, mock_GetOperationStatus, t_transport_class
@@ -1718,7 +1727,8 @@ def test_make_request_will_retry_GetOperationStatus(
         "databricks.sql.thrift_api.TCLIService.TCLIService.Client.GetOperationStatus"
     )
     @patch(
-        "databricks.sql.backend.thrift_backend._retry_policy", new_callable=retry_policy_factory
+        "databricks.sql.backend.thrift_backend._retry_policy",
+        new_callable=retry_policy_factory,
     )
     def test_make_request_will_retry_GetOperationStatus_for_http_error(
         self, mock_retry_policy, mock_gos
@@ -1795,7 +1805,8 @@ def test_make_request_wont_retry_if_error_code_not_429_or_503(
 
     @patch("databricks.sql.auth.thrift_http_client.THttpClient")
     @patch(
-        "databricks.sql.backend.thrift_backend._retry_policy", new_callable=retry_policy_factory
+        "databricks.sql.backend.thrift_backend._retry_policy",
+        new_callable=retry_policy_factory,
     )
     def test_make_request_will_retry_stop_after_attempts_count_if_retryable(
         self, mock_retry_policy, t_transport_class
@@ -1975,7 +1986,12 @@ def test_retry_args_passthrough(self, mock_http_client):
     @patch("thrift.transport.THttpClient.THttpClient")
     def test_retry_args_bounding(self, mock_http_client):
         retry_delay_test_args_and_expected_values = {}
-        for k, (_, _, min, max) in databricks.sql.backend.thrift_backend._retry_policy.items():
+        for k, (
+            _,
+            _,
+            min,
+            max,
+        ) in databricks.sql.backend.thrift_backend._retry_policy.items():
             retry_delay_test_args_and_expected_values[k] = (
                 (min - 1, min),
                 (max + 1, max),
@@ -2171,7 +2187,9 @@ def test_protocol_v3_fails_if_initial_namespace_set(self, tcli_client_class):
         )
 
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
-    @patch("databricks.sql.backend.thrift_backend.ThriftDatabricksClient._handle_execute_response")
+    @patch(
+        "databricks.sql.backend.thrift_backend.ThriftDatabricksClient._handle_execute_response"
+    )
     def test_execute_command_sets_complex_type_fields_correctly(
         self, mock_handle_execute_response, tcli_service_class
     ):