Code import - branch release/SIENTIAPDE-1645
This commit is contained in:
0
tests/schedules/__init__.py
Normal file
0
tests/schedules/__init__.py
Normal file
469
tests/schedules/test_cleanup_schedule.py
Normal file
469
tests/schedules/test_cleanup_schedule.py
Normal file
@@ -0,0 +1,469 @@
|
||||
"""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'
|
||||
Reference in New Issue
Block a user