Apply black formating

Author: neomatamune

cosmotech/coal/cosmotech_api/runner/download.py

@@ -33,7 +33,6 @@
 from cosmotech.orchestrator.utils.translate import T
 
 
-
 def download_runner_data(
     organization_id: str,
     workspace_id: str,

cosmotech/coal/cosmotech_api/twin_data_layer.py

@@ -137,9 +137,7 @@ def generate_query_insert(self) -> str:
         return query
 
 
-def get_dataset_id_from_runner(
-    organization_id: str, workspace_id: str, runner_id: str
-) -> str:
+def get_dataset_id_from_runner(organization_id: str, workspace_id: str, runner_id: str) -> str:
     """
     Get the dataset ID from a runner.
 
@@ -153,7 +151,7 @@ def get_dataset_id_from_runner(
     """
     api_client, _ = get_api_client()
     api_runner = RunnerApi(api_client)
-    
+
     runner_info = api_runner.get_runner(
         organization_id,
         workspace_id,

tests/unit/coal/test_azure/test_adx/test_adx_ingestion_edge_cases.py

@@ -77,15 +77,14 @@ def test_check_ingestion_status_with_logs_and_messages(
 
         # Act
         with patch(
-            "cosmotech.coal.azure.adx.ingestion.SuccessMessage", 
-            return_value=MagicMock(IngestionSourceId=source_id)
+            "cosmotech.coal.azure.adx.ingestion.SuccessMessage", return_value=MagicMock(IngestionSourceId=source_id)
         ):
             result = list(check_ingestion_status(mock_ingest_client, [source_id], logs=True))
 
         # Assert
         assert len(result) == 1
         assert result[0] == (source_id, IngestionStatus.SUCCESS)
-        
+
         # Verify that the message was deleted
         mock_success_queue.delete_message.assert_called_once_with(mock_message)
 
@@ -123,19 +122,16 @@ def test_check_ingestion_status_with_multiple_messages(
 
         # Act
         with patch(
-            "cosmotech.coal.azure.adx.ingestion.SuccessMessage", 
-            side_effect=[
-                MagicMock(IngestionSourceId=source_id1),
-                MagicMock(IngestionSourceId=source_id2)
-            ]
+            "cosmotech.coal.azure.adx.ingestion.SuccessMessage",
+            side_effect=[MagicMock(IngestionSourceId=source_id1), MagicMock(IngestionSourceId=source_id2)],
         ):
             result = list(check_ingestion_status(mock_ingest_client, [source_id1, source_id2], logs=True))
 
         # Assert
         assert len(result) == 2
         assert (source_id1, IngestionStatus.SUCCESS) in result
         assert (source_id2, IngestionStatus.QUEUED) in result or (source_id2, IngestionStatus.SUCCESS) in result
-        
+
         # Verify that at least one message was deleted
         assert mock_success_queue.delete_message.call_count >= 1
 
@@ -166,15 +162,14 @@ def test_check_ingestion_status_with_success_messages_and_logs(
 
         # Act
         with patch(
-            "cosmotech.coal.azure.adx.ingestion.SuccessMessage", 
-            return_value=MagicMock(IngestionSourceId=source_id)
+            "cosmotech.coal.azure.adx.ingestion.SuccessMessage", return_value=MagicMock(IngestionSourceId=source_id)
         ):
             result = list(check_ingestion_status(mock_ingest_client, [source_id], logs=True))
 
         # Assert
         assert len(result) == 1
         assert result[0] == (source_id, IngestionStatus.SUCCESS)
-        
+
         # Verify that the message was deleted
         mock_success_queue.delete_message.assert_called_once_with(mock_success_message)
 
@@ -205,15 +200,14 @@ def test_check_ingestion_status_with_failure_messages_and_logs(
 
         # Act
         with patch(
-            "cosmotech.coal.azure.adx.ingestion.FailureMessage", 
-            return_value=MagicMock(IngestionSourceId=source_id)
+            "cosmotech.coal.azure.adx.ingestion.FailureMessage", return_value=MagicMock(IngestionSourceId=source_id)
         ):
             result = list(check_ingestion_status(mock_ingest_client, [source_id], logs=True))
 
         # Assert
         assert len(result) == 1
         assert result[0] == (source_id, IngestionStatus.FAILURE)
-        
+
         # Verify that the message was deleted
         mock_failure_queue.delete_message.assert_called_once_with(mock_failure_message)
 
@@ -249,18 +243,17 @@ def test_check_ingestion_status_with_logs_and_status_messages(
 
         # Act
         with patch(
-            "cosmotech.coal.azure.adx.ingestion.SuccessMessage", 
-            return_value=MagicMock(IngestionSourceId=source_id)
+            "cosmotech.coal.azure.adx.ingestion.SuccessMessage", return_value=MagicMock(IngestionSourceId=source_id)
         ):
             result = list(check_ingestion_status(mock_ingest_client, [source_id], logs=True))
 
         # Assert
         assert len(result) == 1
         assert result[0] == (source_id, IngestionStatus.SUCCESS)
-        
+
         # Verify that the debug log was called with the correct message
         mock_logger.debug.assert_any_call(T("coal.logs.adx.status_messages").format(success=2, failure=1))
-        
+
         # Verify that the message was deleted
         mock_success_queue.delete_message.assert_called_once_with(mock_success_message1)
 
@@ -291,47 +284,50 @@ def test_check_ingestion_status_with_no_matching_messages(
 
         # Act
         with patch(
-            "cosmotech.coal.azure.adx.ingestion.SuccessMessage", 
-            return_value=MagicMock(IngestionSourceId="different-source-id")
+            "cosmotech.coal.azure.adx.ingestion.SuccessMessage",
+            return_value=MagicMock(IngestionSourceId="different-source-id"),
         ):
             result = list(check_ingestion_status(mock_ingest_client, [source_id], logs=True))
 
         # Assert
         assert len(result) == 1
         assert result[0] == (source_id, IngestionStatus.QUEUED)
-        
+
         # Verify that no messages were deleted
         mock_success_queue.delete_message.assert_not_called()
 
     def test_status_messages_log_line_true(self):
         """Test the specific log line that's not being covered with logs=True."""
         # Import the module directly to access the function
         import cosmotech.coal.azure.adx.ingestion as ingestion_module
-        
+
         # Create mock objects
         mock_logger = MagicMock()
         mock_t = MagicMock()
         mock_format = MagicMock()
         mock_t.return_value = mock_format
         mock_format.format.return_value = "Status message"
-        
+
         # Replace the real objects with mocks
         original_logger = ingestion_module.LOGGER
         original_t = ingestion_module.T
         ingestion_module.LOGGER = mock_logger
         ingestion_module.T = mock_t
-        
+
         try:
             # Create test data
             successes = [1, 2, 3]  # Just need a list with a length
             failures = [1]  # Just need a list with a length
             logs = True
-            
+
             # Call the specific line directly
             if logs:
-                ingestion_module.LOGGER.debug(ingestion_module.T("coal.logs.adx.status_messages").format(
-                    success=len(successes), failure=len(failures)))
-            
+                ingestion_module.LOGGER.debug(
+                    ingestion_module.T("coal.logs.adx.status_messages").format(
+                        success=len(successes), failure=len(failures)
+                    )
+                )
+
             # Verify the mocks were called correctly
             mock_t.assert_called_once_with("coal.logs.adx.status_messages")
             mock_format.format.assert_called_once_with(success=3, failure=1)
@@ -340,33 +336,36 @@ def test_status_messages_log_line_true(self):
             # Restore the original objects
             ingestion_module.LOGGER = original_logger
             ingestion_module.T = original_t
-            
+
     def test_status_messages_log_line_false(self):
         """Test the specific log line that's not being covered with logs=False."""
         # Import the module directly to access the function
         import cosmotech.coal.azure.adx.ingestion as ingestion_module
-        
+
         # Create mock objects
         mock_logger = MagicMock()
         mock_t = MagicMock()
-        
+
         # Replace the real objects with mocks
         original_logger = ingestion_module.LOGGER
         original_t = ingestion_module.T
         ingestion_module.LOGGER = mock_logger
         ingestion_module.T = mock_t
-        
+
         try:
             # Create test data
             successes = [1, 2, 3]  # Just need a list with a length
             failures = [1]  # Just need a list with a length
             logs = False
-            
+
             # Call the specific line directly
             if logs:
-                ingestion_module.LOGGER.debug(ingestion_module.T("coal.logs.adx.status_messages").format(
-                    success=len(successes), failure=len(failures)))
-            
+                ingestion_module.LOGGER.debug(
+                    ingestion_module.T("coal.logs.adx.status_messages").format(
+                        success=len(successes), failure=len(failures)
+                    )
+                )
+
             # Verify the mocks were not called
             mock_t.assert_not_called()
             mock_logger.debug.assert_not_called()
@@ -405,21 +404,20 @@ def test_check_ingestion_status_with_logs_disabled(
 
         # Act
         with patch(
-            "cosmotech.coal.azure.adx.ingestion.SuccessMessage", 
-            return_value=MagicMock(IngestionSourceId=source_id)
+            "cosmotech.coal.azure.adx.ingestion.SuccessMessage", return_value=MagicMock(IngestionSourceId=source_id)
         ):
             result = list(check_ingestion_status(mock_ingest_client, [source_id], logs=False))
 
         # Assert
         assert len(result) == 1
         assert result[0] == (source_id, IngestionStatus.SUCCESS)
-        
+
         # Verify that the debug log was not called with the status messages
         for call_args in mock_logger.debug.call_args_list:
             args, kwargs = call_args
             if len(args) > 0 and isinstance(args[0], str) and "status_messages" in args[0]:
                 assert False, "LOGGER.debug should not be called with status_messages when logs=False"
-        
+
         # Verify that the message was deleted
         mock_success_queue.delete_message.assert_called_once_with(mock_success_message)
 
@@ -443,7 +441,7 @@ def test_check_ingestion_status_with_multiple_queues(
         mock_message = MagicMock()
         mock_message.content = '{"IngestionSourceId": "source-id-multiple-queues"}'
         mock_success_queue2.receive_messages.return_value = [mock_message]
-        
+
         # Set up the success queues
         mock_status_queues.success._get_queues.return_value = [mock_success_queue1, mock_success_queue2]
 
@@ -454,15 +452,14 @@ def test_check_ingestion_status_with_multiple_queues(
 
         # Act
         with patch(
-            "cosmotech.coal.azure.adx.ingestion.SuccessMessage", 
-            return_value=MagicMock(IngestionSourceId=source_id)
+            "cosmotech.coal.azure.adx.ingestion.SuccessMessage", return_value=MagicMock(IngestionSourceId=source_id)
         ):
             result = list(check_ingestion_status(mock_ingest_client, [source_id], logs=True))
 
         # Assert
         assert len(result) == 1
         assert result[0] == (source_id, IngestionStatus.SUCCESS)
-        
+
         # Verify that the message was deleted from the correct queue
         mock_success_queue1.delete_message.assert_not_called()
         mock_success_queue2.delete_message.assert_called_once_with(mock_message)

tests/unit/coal/test_cosmotech_api/test_cosmotech_api_twin_data_layer.py

@@ -170,7 +170,6 @@ def mock_runner_api(self):
         mock_api = MagicMock(spec=RunnerApi)
         return mock_api
 
-
     @pytest.fixture
     def mock_dataset_api(self):
         """Create a mock DatasetApi."""
@@ -196,7 +195,6 @@ def test_get_dataset_id_from_runner(self, mock_get_api_client, mock_runner_api):
             assert result == "dataset-123"
             mock_runner_api.get_runner.assert_called_once_with("org-123", "ws-123", "runner-123")
 
-
     @patch("cosmotech.coal.cosmotech_api.twin_data_layer.get_api_client")
     def test_get_dataset_id_from_runner_no_ids(self, mock_get_api_client, mock_runner_api):
         """Test the get_dataset_id_from_runner function with no IDs."""
@@ -212,7 +210,6 @@ def test_get_dataset_id_from_runner_no_ids(self, mock_get_api_client, mock_runne
             with pytest.raises(ValueError):
                 get_dataset_id_from_runner("org-123", "ws-123", "runner-123")
 
-
     @patch("cosmotech.coal.cosmotech_api.twin_data_layer.get_api_client")
     def test_get_dataset_id_from_runner_no_datasets(self, mock_get_api_client, mock_runner_api):
         """Test the get_dataset_id_from_runner function with no datasets."""

tests/unit/coal/test_cosmotech_api/test_cosmotech_api_twin_data_layer_auth.py

@@ -36,7 +36,7 @@ def mock_api_client(self):
         mock_client.default_headers = {"Default-Header": "value"}
         mock_client.configuration.auth_settings.return_value = {
             "auth1": {"type": "apiKey", "in": "header", "key": "Authorization", "value": "Bearer token"},
-            "auth2": {"type": "basic", "in": "header", "key": "Authorization", "value": "Basic credentials"}
+            "auth2": {"type": "basic", "in": "header", "key": "Authorization", "value": "Basic credentials"},
         }
         return mock_client
 
@@ -59,11 +59,11 @@ def test_send_files_to_tdl_auth_params(
         # Arrange
         mock_api_client = MagicMock()
         mock_api_client.default_headers = {"Default-Header": "value"}
-        
+
         # Set up auth settings with multiple auth types
         auth_settings = {
             "auth1": {"type": "apiKey", "in": "header", "key": "Authorization", "value": "Bearer token"},
-            "auth2": {"type": "basic", "in": "header", "key": "Basic-Auth", "value": "Basic credentials"}
+            "auth2": {"type": "basic", "in": "header", "key": "Basic-Auth", "value": "Basic credentials"},
         }
         mock_api_client.configuration.auth_settings.return_value = auth_settings
 
@@ -96,19 +96,19 @@ def test_send_files_to_tdl_auth_params(
             # Assert
             # Verify that _apply_auth_params was called for each auth type
             assert mock_api_client._apply_auth_params.call_count == len(auth_settings)
-            
+
             # Check that the header was passed to _process_csv_file with auth params applied
             expected_header = {
                 "Accept": "application/json",
                 "Content-Type": "text/csv",
                 "User-Agent": "OpenAPI-Generator/1.0.0/python",
                 "Default-Header": "value",
             }
-            
+
             # Verify that _process_csv_file was called with the expected header
             mock_process_csv_file.assert_called_once()
             actual_header = mock_process_csv_file.call_args[1]["header"]
-            
+
             # Check that the header contains the expected keys
             for key in expected_header:
                 assert key in actual_header

tests/unit/coal/test_cosmotech_api/test_cosmotech_api_twin_data_layer_edge_cases.py

@@ -133,7 +133,7 @@ def test_process_csv_file_with_errors(self, mock_logger, mock_post):
                 _process_csv_file(file_path, query, api_url, organization_id, dataset_id, header)
 
             assert f"Error importing data from {file_path}" in str(excinfo.value)
-            
+
             # Verify that errors were logged
             mock_logger.error.assert_any_call(T("coal.logs.storage.import_errors").format(count=2))
             mock_logger.error.assert_any_call("Error 1")

tests/unit/coal/test_cosmotech_api/test_runner/test_runner_download.py

@@ -22,7 +22,6 @@
 class TestDownloadFunctions:
     """Tests for top-level functions in the download module."""
 
-
     @patch("cosmotech.coal.cosmotech_api.runner.download.get_api_client")
     @patch("cosmotech.coal.cosmotech_api.runner.download.get_runner_data")
     @patch("cosmotech.coal.cosmotech_api.runner.download.format_parameters_list")

tests/unit/coal/test_cosmotech_api/test_runner/test_runner_download_edge_cases.py

@@ -22,7 +22,6 @@
 class TestDownloadEdgeCases:
     """Tests for edge cases in the download module."""
 
-
     @patch("cosmotech.coal.cosmotech_api.runner.download.get_api_client")
     @patch("cosmotech.coal.cosmotech_api.runner.download.DefaultAzureCredential")
     @patch("cosmotech.coal.cosmotech_api.runner.download.get_runner_data")