470 lines
16 KiB
Python
470 lines
16 KiB
Python
"""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()
|
|
handle = AsyncMock()
|
|
handle.delete = AsyncMock()
|
|
schedule = MagicMock()
|
|
schedule.action.task_queue = 'cleanup_files-model-manager-worker-queue'
|
|
schedule.action.execution_timeout = timedelta(hours=1)
|
|
schedule.spec.cron_expressions = ['0 0 * * *']
|
|
schedule.spec.time_zone_name = 'UTC'
|
|
handle.describe = AsyncMock(return_value=MagicMock(schedule=schedule))
|
|
client.get_schedule_handle = MagicMock(return_value=handle)
|
|
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]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_needs_schedule_reconcile_handles_describe_exception(mock_logger, metadata):
|
|
"""Test _needs_schedule_reconcile returns True and logs when describe fails."""
|
|
from model_manager.schedules.cleanup_schedule import _needs_schedule_reconcile
|
|
|
|
handle = AsyncMock()
|
|
handle.describe = AsyncMock(side_effect=RuntimeError('describe failed'))
|
|
|
|
needs_reconcile = await _needs_schedule_reconcile(
|
|
schedule_handle=handle,
|
|
cleanup_task_queue='cleanup_files-model-manager-worker-queue',
|
|
logger=mock_logger,
|
|
metadata=metadata,
|
|
)
|
|
|
|
assert needs_reconcile is True
|
|
mock_logger.custom_error.assert_called_once()
|
|
assert (
|
|
'Error describing cleanup schedule for reconcile'
|
|
in mock_logger.custom_error.call_args[0][0]
|
|
)
|
|
|
|
|
|
# --- create_cleanup_schedule Tests ---
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch.dict(
|
|
'model_manager.schedules.cleanup_schedule.os.environ',
|
|
{
|
|
'RUNTIME': 'model-manager-worker',
|
|
},
|
|
)
|
|
async def test_create_cleanup_schedule_reconciles_when_exists(
|
|
mock_temporal_client, mock_logger, metadata
|
|
):
|
|
"""Test that create_cleanup_schedule recreates schedule when it 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-model-manager-worker-daily'
|
|
|
|
async def mock_list_schedules():
|
|
yield mock_schedule
|
|
|
|
mock_temporal_client.list_schedules.return_value = mock_list_schedules()
|
|
|
|
# Force reconcile by diverging task queue
|
|
mock_temporal_client.get_schedule_handle.return_value.describe.return_value.schedule.action.task_queue = 'different-queue'
|
|
|
|
await create_cleanup_schedule(mock_temporal_client, mock_logger, metadata)
|
|
|
|
# Verify schedule was reconciled via delete + create
|
|
mock_temporal_client.get_schedule_handle.assert_called_once_with(
|
|
'cleanup-files-model-manager-worker-daily'
|
|
)
|
|
mock_temporal_client.get_schedule_handle.return_value.delete.assert_called_once()
|
|
mock_temporal_client.create_schedule.assert_called_once()
|
|
|
|
mock_logger.custom_info.assert_called_once()
|
|
assert 'reconciled successfully' in mock_logger.custom_info.call_args[0][0]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch.dict(
|
|
'model_manager.schedules.cleanup_schedule.os.environ',
|
|
{
|
|
'RUNTIME': 'model-manager-worker',
|
|
},
|
|
)
|
|
async def test_create_cleanup_schedule_noop_when_schedule_is_up_to_date(
|
|
mock_temporal_client, mock_logger, metadata
|
|
):
|
|
"""Test no-op reconcile when existing schedule already matches current config."""
|
|
import model_manager.schedules.cleanup_schedule
|
|
|
|
reload(model_manager.schedules.cleanup_schedule)
|
|
from model_manager.schedules.cleanup_schedule import create_cleanup_schedule
|
|
|
|
mock_schedule = MagicMock()
|
|
mock_schedule.id = 'cleanup-files-model-manager-worker-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)
|
|
|
|
mock_temporal_client.get_schedule_handle.assert_called_once_with(
|
|
'cleanup-files-model-manager-worker-daily'
|
|
)
|
|
mock_temporal_client.get_schedule_handle.return_value.delete.assert_not_called()
|
|
mock_temporal_client.create_schedule.assert_not_called()
|
|
mock_logger.custom_info.assert_called_once()
|
|
assert 'no-op reconcile' in mock_logger.custom_info.call_args[0][0]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch.dict(
|
|
'model_manager.schedules.cleanup_schedule.os.environ',
|
|
{
|
|
'RUNTIME': 'model-manager-worker',
|
|
'CLEANUP_CRON': '0 2 * * *',
|
|
'CLEANUP_TIMEZONE': 'America/Sao_Paulo',
|
|
'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 == 'cleanup-files-model-manager-worker-daily'
|
|
assert schedule_obj.action.workflow == 'cleanup_files'
|
|
assert schedule_obj.action.task_queue == 'cleanup_files-model-manager-worker-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',
|
|
{
|
|
'RUNTIME': 'model-manager-worker',
|
|
},
|
|
)
|
|
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_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_files-model-manager-worker-queue'
|
|
assert schedule_obj.action.execution_timeout == timedelta(hours=1) # Default 1 hour
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch.dict(
|
|
'model_manager.schedules.cleanup_schedule.os.environ',
|
|
{
|
|
'RUNTIME': 'model-manager-worker',
|
|
},
|
|
)
|
|
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 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 = 'cleanup-files-scheduled-cleanup-files-model-manager-worker-daily'
|
|
assert schedule_obj.action.id == expected_workflow_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch.dict(
|
|
'model_manager.schedules.cleanup_schedule.os.environ',
|
|
{
|
|
'RUNTIME': 'model-manager-worker',
|
|
},
|
|
)
|
|
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',
|
|
{
|
|
'RUNTIME': 'model-manager-worker',
|
|
'CLEANUP_CRON': '30 3 * * 1',
|
|
'CLEANUP_TIMEZONE': 'Europe/London',
|
|
'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_TIMEZONE,
|
|
build_cleanup_schedule_id,
|
|
)
|
|
|
|
assert (
|
|
build_cleanup_schedule_id('model-manager-worker')
|
|
== 'cleanup-files-model-manager-worker-daily'
|
|
)
|
|
assert CLEANUP_CRON == '30 3 * * 1'
|
|
assert CLEANUP_TIMEZONE == 'Europe/London'
|
|
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 [
|
|
'RUNTIME',
|
|
'CLEANUP_CRON',
|
|
'CLEANUP_TIMEZONE',
|
|
'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_TIMEZONE,
|
|
build_cleanup_schedule_id,
|
|
)
|
|
|
|
assert build_cleanup_schedule_id(None) == 'cleanup-files-single-daily'
|
|
assert CLEANUP_CRON == '0 0 * * *'
|
|
assert CLEANUP_TIMEZONE == 'UTC'
|
|
assert CLEANUP_EXECUTION_TIMEOUT_HOURS == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_cleanup_schedule_uses_single_runtime_when_runtime_missing(
|
|
mock_temporal_client, mock_logger, metadata
|
|
):
|
|
"""Test create_cleanup_schedule uses single runtime fallback."""
|
|
import model_manager.schedules.cleanup_schedule
|
|
|
|
os.environ.pop('RUNTIME', None)
|
|
reload(model_manager.schedules.cleanup_schedule)
|
|
from model_manager.schedules.cleanup_schedule import create_cleanup_schedule
|
|
|
|
async def mock_list_schedules():
|
|
return
|
|
yield
|
|
|
|
mock_temporal_client.list_schedules.return_value = mock_list_schedules()
|
|
|
|
await create_cleanup_schedule(mock_temporal_client, mock_logger, metadata)
|
|
|
|
call_args = mock_temporal_client.create_schedule.call_args
|
|
schedule_obj = call_args[0][1]
|
|
assert schedule_obj.action.task_queue == 'cleanup_files-single-queue'
|