gitea_admin/sientia-dataops-model-manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1350: Implement scheduled cleanup workflow using Temporal schedules. Adds schedule creation to worker startup and configures environment variables.

Browse Source
This commit is contained in:
Bruno Domingues
2025-11-25 18:12:53 -03:00
parent 9ecac9c8fd
commit b028dc1b7d
10 changed files with 140 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
tests/worker/test_worker.py
View File

@@ -335,6 +335,7 @@ async def test_main_handles_exception(
@pytest.mark.asyncio
@patch('model_manager.worker.worker.create_cleanup_schedule')
@patch('model_manager.worker.worker.Worker')
@patch('model_manager.worker.worker.client.Client')
@patch('model_manager.worker.worker.Runtime')
@@ -360,11 +361,14 @@ async def test_main_temporal_client_configuration(
mock_runtime_class,
mock_client_class,
mock_worker_class,
mock_create_cleanup_schedule,
mock_logger,
):
"""Test that Temporal client is configured correctly."""
from model_manager.worker.worker import main
mock_create_cleanup_schedule.return_value = AsyncMock()
with patch.dict(
os.environ,
{'TEMPORAL_HOST': 'temporal.example.com:7233', 'TEMPORAL_NAMESPACE': 'production'},
@@ -410,11 +414,12 @@ async def test_main_temporal_client_configuration(
target_host='temporal.example.com:7233',
namespace='production',
runtime=mock_runtime,
tls=False,
tls=True,
)
@pytest.mark.asyncio
@patch('model_manager.worker.worker.create_cleanup_schedule')
@patch('model_manager.worker.worker.Worker')
@patch('model_manager.worker.worker.client.Client')
@patch('model_manager.worker.worker.Runtime')
@@ -440,12 +445,15 @@ async def test_main_worker_configuration(
mock_runtime_class,
mock_client_class,
mock_worker_class,
mock_create_cleanup_schedule,
mock_env_vars,
mock_logger,
):
"""Test that Temporal worker is configured with correct parameters."""
from model_manager.worker.worker import main
mock_create_cleanup_schedule.return_value = AsyncMock()
# Setup mocks
mock_get_logger.return_value = mock_logger
mock_build_mongodb.return_value = {
@@ -488,10 +496,10 @@ async def test_main_worker_configuration(
# Verify Worker was created with correct configuration
assert mock_worker_class.call_count == 2
# Primeira chamada: worker de treinamento (train_model-queue)
# 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-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_activities'] == 10
assert train_call_args[1]['max_concurrent_local_activities'] == 10
@@ -503,10 +511,10 @@ async def test_main_worker_configuration(
assert mock_activities.train_model in train_activities_list
assert mock_activities.cleanup_resources in train_activities_list
# Segunda chamada: worker de cleanup (cleanup-queue)
# 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-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_activities'] == 20
assert cleanup_call_args[1]['max_concurrent_local_activities'] == 20