From bc917e97d84f61a81ab4aa4dedc7432969d12453 Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Tue, 25 Nov 2025 19:21:55 -0300 Subject: [PATCH] SIENTIAPDE-1350: Add cleanup schedule tests and configure worker with task queues from env vars --- tests/schedules/test_cleanup_schedule.py | 362 +++++++++++++++++++++++ tests/worker/test_worker.py | 115 +++---- 2 files changed, 423 insertions(+), 54 deletions(-) create mode 100644 tests/schedules/test_cleanup_schedule.py diff --git a/tests/schedules/test_cleanup_schedule.py b/tests/schedules/test_cleanup_schedule.py new file mode 100644 index 0000000..72ba651 --- /dev/null +++ b/tests/schedules/test_cleanup_schedule.py @@ -0,0 +1,362 @@ +"""Tests for cleanup schedule management.""" + +import os +from datetime import timedelta +from importlib import reload +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + + +@pytest.fixture +def mock_temporal_client(): + """Fixture for a mock Temporal client.""" + client = AsyncMock() + client.list_schedules = AsyncMock() + client.create_schedule = AsyncMock() + return client + + +@pytest.fixture +def mock_logger(): + """Fixture for a mock Sientia logger.""" + logger = MagicMock() + logger.custom_info = MagicMock() + logger.custom_error = MagicMock() + return logger + + +@pytest.fixture +def metadata(): + """Fixture for metadata dict.""" + return {'pod_id': 'test-pod', 'project_name': 'test-project'} + + +# --- schedule_exists Tests --- + + +@pytest.mark.asyncio +async def test_schedule_exists_returns_true_when_schedule_found( + mock_temporal_client, mock_logger, metadata +): + """Test that schedule_exists returns True when schedule is found.""" + from model_manager.schedules.cleanup_schedule import schedule_exists + + # Mock schedule list with matching schedule + mock_schedule = MagicMock() + mock_schedule.id = 'test-schedule-id' + + async def mock_list_schedules(): + yield mock_schedule + + mock_temporal_client.list_schedules.return_value = mock_list_schedules() + + result = await schedule_exists(mock_temporal_client, 'test-schedule-id', mock_logger, metadata) + + assert result is True + mock_temporal_client.list_schedules.assert_called_once() + + +@pytest.mark.asyncio +async def test_schedule_exists_returns_false_when_schedule_not_found( + mock_temporal_client, mock_logger, metadata +): + """Test that schedule_exists returns False when schedule is not found.""" + from model_manager.schedules.cleanup_schedule import schedule_exists + + # Mock empty schedule list + async def mock_list_schedules(): + return + yield # Make it an async generator + + mock_temporal_client.list_schedules.return_value = mock_list_schedules() + + result = await schedule_exists( + mock_temporal_client, 'nonexistent-schedule', mock_logger, metadata + ) + + assert result is False + mock_temporal_client.list_schedules.assert_called_once() + + +@pytest.mark.asyncio +async def test_schedule_exists_returns_false_when_different_schedule_found( + mock_temporal_client, mock_logger, metadata +): + """Test that schedule_exists returns False when only different schedules exist.""" + from model_manager.schedules.cleanup_schedule import schedule_exists + + # Mock schedule list with non-matching schedule + mock_schedule = MagicMock() + mock_schedule.id = 'different-schedule-id' + + async def mock_list_schedules(): + yield mock_schedule + + mock_temporal_client.list_schedules.return_value = mock_list_schedules() + + result = await schedule_exists(mock_temporal_client, 'test-schedule-id', mock_logger, metadata) + + assert result is False + mock_temporal_client.list_schedules.assert_called_once() + + +@pytest.mark.asyncio +async def test_schedule_exists_handles_exception(mock_temporal_client, mock_logger, metadata): + """Test that schedule_exists handles exceptions gracefully.""" + from model_manager.schedules.cleanup_schedule import schedule_exists + + # Mock list_schedules to raise an exception + mock_temporal_client.list_schedules.side_effect = Exception('Connection error') + + result = await schedule_exists(mock_temporal_client, 'test-schedule-id', mock_logger, metadata) + + assert result is False + mock_logger.custom_error.assert_called_once() + assert 'Error checking if schedule exists' in mock_logger.custom_error.call_args[0][0] + + +# --- create_cleanup_schedule Tests --- + + +@pytest.mark.asyncio +async def test_create_cleanup_schedule_skips_when_exists( + mock_temporal_client, mock_logger, metadata +): + """Test that create_cleanup_schedule skips creation when schedule already exists.""" + import model_manager.schedules.cleanup_schedule + + reload(model_manager.schedules.cleanup_schedule) + from model_manager.schedules.cleanup_schedule import create_cleanup_schedule + + # Mock schedule already exists + mock_schedule = MagicMock() + mock_schedule.id = 'cleanup-files-daily' + + async def mock_list_schedules(): + yield mock_schedule + + mock_temporal_client.list_schedules.return_value = mock_list_schedules() + + await create_cleanup_schedule(mock_temporal_client, mock_logger, metadata) + + # Verify schedule creation was NOT called + mock_temporal_client.create_schedule.assert_not_called() + + # Verify info log was called + mock_logger.custom_info.assert_called_once() + assert 'already configured' in mock_logger.custom_info.call_args[0][0] + + +@pytest.mark.asyncio +@patch.dict( + 'model_manager.schedules.cleanup_schedule.os.environ', + { + 'CLEANUP_SCHEDULE_ID': 'test-cleanup-schedule', + 'CLEANUP_CRON': '0 2 * * *', + 'CLEANUP_TIMEZONE': 'America/Sao_Paulo', + 'CLEANUP_TASK_QUEUE': 'test-cleanup-queue', + 'CLEANUP_EXECUTION_TIMEOUT_HOURS': '2', + }, +) +async def test_create_cleanup_schedule_creates_with_custom_config( + mock_temporal_client, mock_logger, metadata +): + """Test that create_cleanup_schedule creates schedule with custom configuration.""" + import model_manager.schedules.cleanup_schedule + + reload(model_manager.schedules.cleanup_schedule) + from model_manager.schedules.cleanup_schedule import create_cleanup_schedule + + # Mock schedule does not exist (empty list) + async def mock_list_schedules(): + return + yield # Make it an async generator + + mock_temporal_client.list_schedules.return_value = mock_list_schedules() + + await create_cleanup_schedule(mock_temporal_client, mock_logger, metadata) + + # Verify schedule creation was called + mock_temporal_client.create_schedule.assert_called_once() + + # Verify schedule parameters + call_args = mock_temporal_client.create_schedule.call_args + schedule_id = call_args[0][0] + schedule_obj = call_args[0][1] + + assert schedule_id == 'test-cleanup-schedule' + assert schedule_obj.action.workflow == 'cleanup_files' + assert schedule_obj.action.task_queue == 'test-cleanup-queue' + assert schedule_obj.action.execution_timeout == timedelta(hours=2) + assert schedule_obj.spec.cron_expressions == ['0 2 * * *'] + assert schedule_obj.spec.time_zone_name == 'America/Sao_Paulo' + + # Verify success log was called + assert mock_logger.custom_info.call_count == 1 + assert 'created successfully' in mock_logger.custom_info.call_args[0][0] + + +@pytest.mark.asyncio +@patch.dict( + 'model_manager.schedules.cleanup_schedule.os.environ', + { + 'CLEANUP_SCHEDULE_ID': 'default-schedule', + }, +) +async def test_create_cleanup_schedule_uses_defaults(mock_temporal_client, mock_logger, metadata): + """Test that create_cleanup_schedule uses default values when env vars not set.""" + import model_manager.schedules.cleanup_schedule + + # Remove optional env vars to test defaults + for key in [ + 'CLEANUP_CRON', + 'CLEANUP_TIMEZONE', + 'CLEANUP_TASK_QUEUE', + 'CLEANUP_EXECUTION_TIMEOUT_HOURS', + ]: + os.environ.pop(key, None) + + reload(model_manager.schedules.cleanup_schedule) + from model_manager.schedules.cleanup_schedule import create_cleanup_schedule + + # Mock schedule does not exist (empty list) + async def mock_list_schedules(): + return + yield # Make it an async generator + + mock_temporal_client.list_schedules.return_value = mock_list_schedules() + + await create_cleanup_schedule(mock_temporal_client, mock_logger, metadata) + + # Verify schedule creation was called + mock_temporal_client.create_schedule.assert_called_once() + + # Verify default parameters + call_args = mock_temporal_client.create_schedule.call_args + schedule_obj = call_args[0][1] + + assert schedule_obj.spec.cron_expressions == ['0 0 * * *'] # Default midnight + assert schedule_obj.spec.time_zone_name == 'UTC' # Default UTC + assert schedule_obj.action.task_queue == 'cleanup_queue' # Default queue + assert schedule_obj.action.execution_timeout == timedelta(hours=1) # Default 1 hour + + +@pytest.mark.asyncio +async def test_create_cleanup_schedule_workflow_id_format( + mock_temporal_client, mock_logger, metadata +): + """Test that workflow ID is correctly formatted with schedule ID.""" + import model_manager.schedules.cleanup_schedule + + reload(model_manager.schedules.cleanup_schedule) + from model_manager.schedules.cleanup_schedule import ( + SCHEDULE_ID, + create_cleanup_schedule, + ) + + # Mock schedule does not exist (empty list) + async def mock_list_schedules(): + return + yield # Make it an async generator + + mock_temporal_client.list_schedules.return_value = mock_list_schedules() + + await create_cleanup_schedule(mock_temporal_client, mock_logger, metadata) + + # Verify workflow ID format + call_args = mock_temporal_client.create_schedule.call_args + schedule_obj = call_args[0][1] + + expected_workflow_id = f'cleanup-files-scheduled-{SCHEDULE_ID}' + assert schedule_obj.action.id == expected_workflow_id + + +@pytest.mark.asyncio +async def test_create_cleanup_schedule_empty_workflow_args( + mock_temporal_client, mock_logger, metadata +): + """Test that workflow is created with empty args (uses env defaults).""" + import model_manager.schedules.cleanup_schedule + + reload(model_manager.schedules.cleanup_schedule) + from model_manager.schedules.cleanup_schedule import create_cleanup_schedule + + # Mock schedule does not exist (empty list) + async def mock_list_schedules(): + return + yield # Make it an async generator + + mock_temporal_client.list_schedules.return_value = mock_list_schedules() + + await create_cleanup_schedule(mock_temporal_client, mock_logger, metadata) + + # Verify workflow args are empty (it's a list with one empty dict) + call_args = mock_temporal_client.create_schedule.call_args + schedule_obj = call_args[0][1] + + # The args are passed as positional args, so it's a list with one element + assert schedule_obj.action.args == [{}] + + +# --- Environment Variable Configuration Tests --- + + +@patch.dict( + 'model_manager.schedules.cleanup_schedule.os.environ', + { + 'CLEANUP_SCHEDULE_ID': 'custom-id', + 'CLEANUP_CRON': '30 3 * * 1', + 'CLEANUP_TIMEZONE': 'Europe/London', + 'CLEANUP_TASK_QUEUE': 'custom-queue', + 'CLEANUP_EXECUTION_TIMEOUT_HOURS': '3', + }, +) +def test_environment_variables_loaded_correctly(): + """Test that environment variables are loaded correctly.""" + import model_manager.schedules.cleanup_schedule + + reload(model_manager.schedules.cleanup_schedule) + from model_manager.schedules.cleanup_schedule import ( + CLEANUP_CRON, + CLEANUP_EXECUTION_TIMEOUT_HOURS, + CLEANUP_TASK_QUEUE, + CLEANUP_TIMEZONE, + SCHEDULE_ID, + ) + + assert SCHEDULE_ID == 'custom-id' + assert CLEANUP_CRON == '30 3 * * 1' + assert CLEANUP_TIMEZONE == 'Europe/London' + assert CLEANUP_TASK_QUEUE == 'custom-queue' + assert CLEANUP_EXECUTION_TIMEOUT_HOURS == 3 + + +def test_environment_variables_use_defaults_when_not_set(): + """Test that default values are used when environment variables are not set.""" + import model_manager.schedules.cleanup_schedule + + # Remove all env vars + for key in [ + 'CLEANUP_SCHEDULE_ID', + 'CLEANUP_CRON', + 'CLEANUP_TIMEZONE', + 'CLEANUP_TASK_QUEUE', + 'CLEANUP_EXECUTION_TIMEOUT_HOURS', + ]: + os.environ.pop(key, None) + + reload(model_manager.schedules.cleanup_schedule) + from model_manager.schedules.cleanup_schedule import ( + CLEANUP_CRON, + CLEANUP_EXECUTION_TIMEOUT_HOURS, + CLEANUP_TASK_QUEUE, + CLEANUP_TIMEZONE, + SCHEDULE_ID, + ) + + assert SCHEDULE_ID == 'cleanup-files-daily' + assert CLEANUP_CRON == '0 0 * * *' + assert CLEANUP_TIMEZONE == 'UTC' + assert CLEANUP_TASK_QUEUE == 'cleanup_queue' + assert CLEANUP_EXECUTION_TIMEOUT_HOURS == 1 diff --git a/tests/worker/test_worker.py b/tests/worker/test_worker.py index 1b81da1..4980b0d 100644 --- a/tests/worker/test_worker.py +++ b/tests/worker/test_worker.py @@ -18,6 +18,8 @@ def mock_env_vars(): 'TEMPORAL_HOST': 'localhost:7233', 'TEMPORAL_NAMESPACE': 'test-namespace', 'PROJECT_NAME': 'test-project', + 'TRAIN_TASK_QUEUE': 'train_model-local_queue', + 'CLEANUP_TASK_QUEUE': 'cleanup-local_queue', } with patch.dict(os.environ, env_vars, clear=False): @@ -454,71 +456,76 @@ async def test_main_worker_configuration( mock_create_cleanup_schedule.return_value = AsyncMock() - # Setup mocks - mock_get_logger.return_value = mock_logger - mock_build_mongodb.return_value = { - 'connection_string': 'mongodb://test', - 'database_name': 'test_db', - 'uri': 'localhost:27018', - } - mock_build_postgres.return_value = {} - mock_build_mlflow.return_value = {} - mock_build_minio.return_value = {} + # Patch the task queue constants directly + with ( + patch('model_manager.worker.worker.TRAIN_TASK_QUEUE', 'train_model-local_queue'), + patch('model_manager.worker.worker.CLEANUP_TASK_QUEUE', 'cleanup-local_queue'), + ): + # Setup mocks + mock_get_logger.return_value = mock_logger + mock_build_mongodb.return_value = { + 'connection_string': 'mongodb://test', + 'database_name': 'test_db', + 'uri': 'localhost:27018', + } + mock_build_postgres.return_value = {} + mock_build_mlflow.return_value = {} + mock_build_minio.return_value = {} - mock_notification_handler = Mock() - mock_notification_handler_class.return_value = mock_notification_handler + mock_notification_handler = Mock() + mock_notification_handler_class.return_value = mock_notification_handler - mock_activities = AsyncMock() - mock_activities.update_experiment_run = Mock() - mock_activities.validate_train_params = Mock() - mock_activities.train_model = Mock() - mock_activities.cleanup_resources = Mock() - mock_activities.shutdown = AsyncMock() - mock_activities_class.return_value = mock_activities + mock_activities = AsyncMock() + mock_activities.update_experiment_run = Mock() + mock_activities.validate_train_params = Mock() + mock_activities.train_model = Mock() + mock_activities.cleanup_resources = Mock() + mock_activities.shutdown = AsyncMock() + mock_activities_class.return_value = mock_activities - mock_runtime = Mock() - mock_runtime_class.return_value = mock_runtime + mock_runtime = Mock() + mock_runtime_class.return_value = mock_runtime - mock_client_instance = AsyncMock() - mock_client_class.connect = AsyncMock(return_value=mock_client_instance) + mock_client_instance = AsyncMock() + mock_client_class.connect = AsyncMock(return_value=mock_client_instance) - mock_worker_instance = Mock() - mock_worker_instance.run = AsyncMock(side_effect=asyncio.CancelledError()) - mock_worker_class.return_value = mock_worker_instance + mock_worker_instance = Mock() + mock_worker_instance.run = AsyncMock(side_effect=asyncio.CancelledError()) + mock_worker_class.return_value = mock_worker_instance - mock_app_up = Mock() - mock_metrics.APP_UP.labels.return_value = mock_app_up + mock_app_up = Mock() + mock_metrics.APP_UP.labels.return_value = mock_app_up - # Run main() - with pytest.raises(SystemExit): - await main() + # Run main() + with pytest.raises(SystemExit): + await main() - # Verify Worker was created with correct configuration - assert mock_worker_class.call_count == 2 + # Verify Worker was created with correct configuration + assert mock_worker_class.call_count == 2 - # Primeira chamada: worker de treinamento (train_model-local_queue) - train_call_args = mock_worker_class.call_args_list[0] - assert train_call_args[0][0] == mock_client_instance # temporal_client - assert train_call_args[1]['task_queue'] == 'train_model-local_queue' - assert train_call_args[1]['max_concurrent_workflow_tasks'] == 10 - assert train_call_args[1]['max_concurrent_activities'] == 10 - assert train_call_args[1]['max_concurrent_local_activities'] == 10 - assert train_call_args[1]['max_cached_workflows'] == 100 + # Primeira chamada: worker de treinamento (train_model-local_queue) + train_call_args = mock_worker_class.call_args_list[0] + assert train_call_args[0][0] == mock_client_instance # temporal_client + assert train_call_args[1]['task_queue'] == 'train_model-local_queue' + assert train_call_args[1]['max_concurrent_workflow_tasks'] == 10 + assert train_call_args[1]['max_concurrent_activities'] == 10 + assert train_call_args[1]['max_concurrent_local_activities'] == 10 + assert train_call_args[1]['max_cached_workflows'] == 100 - train_activities_list = train_call_args[1]['activities'] - assert mock_activities.update_experiment_run in train_activities_list - assert mock_activities.validate_train_params in train_activities_list - assert mock_activities.train_model in train_activities_list - assert mock_activities.cleanup_resources in train_activities_list + train_activities_list = train_call_args[1]['activities'] + assert mock_activities.update_experiment_run in train_activities_list + assert mock_activities.validate_train_params in train_activities_list + assert mock_activities.train_model in train_activities_list + assert mock_activities.cleanup_resources in train_activities_list - # Segunda chamada: worker de cleanup (cleanup-local_queue) - cleanup_call_args = mock_worker_class.call_args_list[1] - assert cleanup_call_args[0][0] == mock_client_instance # temporal_client - assert cleanup_call_args[1]['task_queue'] == 'cleanup-local_queue' - assert cleanup_call_args[1]['max_concurrent_workflow_tasks'] == 20 - assert cleanup_call_args[1]['max_concurrent_activities'] == 20 - assert cleanup_call_args[1]['max_concurrent_local_activities'] == 20 - assert cleanup_call_args[1]['max_cached_workflows'] == 100 + # Segunda chamada: worker de cleanup (cleanup-local_queue) + cleanup_call_args = mock_worker_class.call_args_list[1] + assert cleanup_call_args[0][0] == mock_client_instance # temporal_client + assert cleanup_call_args[1]['task_queue'] == 'cleanup-local_queue' + assert cleanup_call_args[1]['max_concurrent_workflow_tasks'] == 20 + assert cleanup_call_args[1]['max_concurrent_activities'] == 20 + assert cleanup_call_args[1]['max_concurrent_local_activities'] == 20 + assert cleanup_call_args[1]['max_cached_workflows'] == 100 @patch('model_manager.worker.worker.asyncio.run')