SIENTIAPDE-1350: Add unit tests for the CleanupFiles workflow, covering both input bucket and default bucket scenarios.
This commit is contained in:
76
tests/workflows/test_cleanup_files.py
Normal file
76
tests/workflows/test_cleanup_files.py
Normal file
@@ -0,0 +1,76 @@
|
||||
"""Unit tests for the CleanupFiles workflow."""
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('model_manager.workflows.cleanup_files.workflow')
|
||||
@patch('model_manager.workflows.cleanup_files.POD_ID', 'temporal-pod')
|
||||
async def test_cleanup_files_workflow_with_input_bucket(mock_workflow_module):
|
||||
"""Test the CleanupFiles workflow when bucket_name is provided in the input."""
|
||||
from model_manager.workflows.cleanup_files import CleanupFiles
|
||||
|
||||
# Mock execute_activity_method
|
||||
mock_workflow_module.execute_activity_method = AsyncMock()
|
||||
|
||||
# Instantiate and run the workflow
|
||||
workflow_instance = CleanupFiles()
|
||||
await workflow_instance.run({'bucket_name': 'input-bucket'})
|
||||
|
||||
# Verify that the activities were called with the correct parameters
|
||||
calls = mock_workflow_module.execute_activity_method.call_args_list
|
||||
assert len(calls) == 2
|
||||
|
||||
# Check cleanup_minio_files call
|
||||
minio_call_args = calls[0][0][1]
|
||||
assert minio_call_args['bucket_name'] == 'input-bucket'
|
||||
assert minio_call_args['metadata'] == {
|
||||
'pod_id': 'temporal-pod',
|
||||
'workflow_name': 'cleanup_files',
|
||||
}
|
||||
|
||||
# Check cleanup_temp_directories call
|
||||
local_call_args = calls[1][0][1]
|
||||
assert local_call_args['temp_path'] == 'model_manager/reports/temp'
|
||||
assert local_call_args['metadata'] == {
|
||||
'pod_id': 'temporal-pod',
|
||||
'workflow_name': 'cleanup_files',
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('model_manager.workflows.cleanup_files.workflow')
|
||||
@patch('model_manager.workflows.cleanup_files.POD_ID', 'temporal-pod')
|
||||
@patch('model_manager.workflows.cleanup_files.DEFAULT_CLEANUP_BUCKET', 'env-var-bucket')
|
||||
async def test_cleanup_files_workflow_with_default_bucket(mock_workflow_module):
|
||||
"""Test the CleanupFiles workflow when using the default bucket from environment variables."""
|
||||
from model_manager.workflows.cleanup_files import CleanupFiles
|
||||
|
||||
# Mock execute_activity_method
|
||||
mock_workflow_module.execute_activity_method = AsyncMock()
|
||||
|
||||
# Instantiate and run the workflow
|
||||
workflow_instance = CleanupFiles()
|
||||
await workflow_instance.run({}) # Empty input
|
||||
|
||||
# Verify that the activities were called
|
||||
calls = mock_workflow_module.execute_activity_method.call_args_list
|
||||
assert len(calls) == 2
|
||||
|
||||
# Check cleanup_minio_files call
|
||||
minio_call_args = calls[0][0][1]
|
||||
assert minio_call_args['bucket_name'] == 'env-var-bucket'
|
||||
assert minio_call_args['metadata'] == {
|
||||
'pod_id': 'temporal-pod',
|
||||
'workflow_name': 'cleanup_files',
|
||||
}
|
||||
|
||||
# Check cleanup_temp_directories call
|
||||
local_call_args = calls[1][0][1]
|
||||
assert local_call_args['temp_path'] == 'model_manager/reports/temp'
|
||||
assert local_call_args['metadata'] == {
|
||||
'pod_id': 'temporal-pod',
|
||||
'workflow_name': 'cleanup_files',
|
||||
}
|
||||
Reference in New Issue
Block a user