SIENTIAPDE-1350: Add cleanup schedule tests and configure worker with task queues from env vars

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

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
+    # Patch the task queue constants directly
-    mock_get_logger.return_value = mock_logger
+    with (
-    mock_build_mongodb.return_value = {
+        patch('model_manager.worker.worker.TRAIN_TASK_QUEUE', 'train_model-local_queue'),
-        'connection_string': 'mongodb://test',
+        patch('model_manager.worker.worker.CLEANUP_TASK_QUEUE', 'cleanup-local_queue'),
-        'database_name': 'test_db',
+    ):
-        'uri': 'localhost:27018',
+        # Setup mocks
-    }
+        mock_get_logger.return_value = mock_logger
-    mock_build_postgres.return_value = {}
+        mock_build_mongodb.return_value = {
-    mock_build_mlflow.return_value = {}
+            'connection_string': 'mongodb://test',
-    mock_build_minio.return_value = {}
+            '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 = Mock()
-    mock_notification_handler_class.return_value = mock_notification_handler
+        mock_notification_handler_class.return_value = mock_notification_handler
 
-    mock_activities = AsyncMock()
+        mock_activities = AsyncMock()
-    mock_activities.update_experiment_run = Mock()
+        mock_activities.update_experiment_run = Mock()
-    mock_activities.validate_train_params = Mock()
+        mock_activities.validate_train_params = Mock()
-    mock_activities.train_model = Mock()
+        mock_activities.train_model = Mock()
-    mock_activities.cleanup_resources = Mock()
+        mock_activities.cleanup_resources = Mock()
-    mock_activities.shutdown = AsyncMock()
+        mock_activities.shutdown = AsyncMock()
-    mock_activities_class.return_value = mock_activities
+        mock_activities_class.return_value = mock_activities
 
-    mock_runtime = Mock()
+        mock_runtime = Mock()
-    mock_runtime_class.return_value = mock_runtime
+        mock_runtime_class.return_value = mock_runtime
 
-    mock_client_instance = AsyncMock()
+        mock_client_instance = AsyncMock()
-    mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
+        mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
 
-    mock_worker_instance = Mock()
+        mock_worker_instance = Mock()
-    mock_worker_instance.run = AsyncMock(side_effect=asyncio.CancelledError())
+        mock_worker_instance.run = AsyncMock(side_effect=asyncio.CancelledError())
-    mock_worker_class.return_value = mock_worker_instance
+        mock_worker_class.return_value = mock_worker_instance
 
-    mock_app_up = Mock()
+        mock_app_up = Mock()
-    mock_metrics.APP_UP.labels.return_value = mock_app_up
+        mock_metrics.APP_UP.labels.return_value = mock_app_up
 
-    # Run main()
+        # Run main()
-    with pytest.raises(SystemExit):
+        with pytest.raises(SystemExit):
-        await main()
+            await main()
 
-    # Verify Worker was created with correct configuration
+        # Verify Worker was created with correct configuration
-    assert mock_worker_class.call_count == 2
+        assert mock_worker_class.call_count == 2
 
-    # Primeira chamada: worker de treinamento (train_model-local_queue)
+        # Primeira chamada: worker de treinamento (train_model-local_queue)
-    train_call_args = mock_worker_class.call_args_list[0]
+        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[0][0] == mock_client_instance  # temporal_client
-    assert train_call_args[1]['task_queue'] == 'train_model-local_queue'
+        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_workflow_tasks'] == 10
-    assert train_call_args[1]['max_concurrent_activities'] == 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_concurrent_local_activities'] == 10
-    assert train_call_args[1]['max_cached_workflows'] == 100
+        assert train_call_args[1]['max_cached_workflows'] == 100
 
-    train_activities_list = train_call_args[1]['activities']
+        train_activities_list = train_call_args[1]['activities']
-    assert mock_activities.update_experiment_run in train_activities_list
+        assert mock_activities.update_experiment_run in train_activities_list
-    assert mock_activities.validate_train_params 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.train_model in train_activities_list
-    assert mock_activities.cleanup_resources in train_activities_list
+        assert mock_activities.cleanup_resources in train_activities_list
 
-    # Segunda chamada: worker de cleanup (cleanup-local_queue)
+        # Segunda chamada: worker de cleanup (cleanup-local_queue)
-    cleanup_call_args = mock_worker_class.call_args_list[1]
+        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[0][0] == mock_client_instance  # temporal_client
-    assert cleanup_call_args[1]['task_queue'] == 'cleanup-local_queue'
+        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_workflow_tasks'] == 20
-    assert cleanup_call_args[1]['max_concurrent_activities'] == 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_concurrent_local_activities'] == 20
-    assert cleanup_call_args[1]['max_cached_workflows'] == 100
+        assert cleanup_call_args[1]['max_cached_workflows'] == 100
 
 
 @patch('model_manager.worker.worker.asyncio.run')