85 lines
2.9 KiB
Python
85 lines
2.9 KiB
Python
"""Unit tests for local worker factory."""
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
|
|
def test_build_queue_name_without_runtime_uses_default_suffix():
|
|
from model_manager.worker.prepare_worker import build_queue_name
|
|
|
|
assert build_queue_name('TrainModel') == 'train_model-queue'
|
|
|
|
|
|
def test_prepare_worker_train_queue_uses_train_limits():
|
|
from model_manager.worker.prepare_worker import prepare_worker
|
|
from model_manager.workflows.train_model import TrainModel
|
|
|
|
fake_worker = MagicMock()
|
|
fake_client = MagicMock()
|
|
fake_logger = MagicMock()
|
|
|
|
with patch(
|
|
'model_manager.worker.prepare_worker.Worker', return_value=fake_worker
|
|
) as worker_class:
|
|
with patch.dict(
|
|
'os.environ',
|
|
{
|
|
'TRAINMODEL_ACTIVITY_EXECUTOR_MAX_WORKERS': '3',
|
|
'TRAINMODEL_MAX_CONCURRENT_ACTIVITIES': '6',
|
|
'TRAINMODEL_MAX_CONCURRENT_WORKFLOW_TASKS': '10',
|
|
},
|
|
clear=False,
|
|
):
|
|
worker = prepare_worker(
|
|
main_workflow=TrainModel,
|
|
other_workflows=[],
|
|
activities=[],
|
|
temporal_client=fake_client,
|
|
logger=fake_logger,
|
|
runtime='model-manager-worker',
|
|
)
|
|
|
|
assert worker is fake_worker
|
|
worker_class.assert_called_once()
|
|
kwargs = worker_class.call_args.kwargs
|
|
assert kwargs['task_queue'] == 'train_model-model-manager-worker-queue'
|
|
assert kwargs['max_concurrent_activities'] == 6
|
|
assert kwargs['max_concurrent_workflow_tasks'] == 10
|
|
assert kwargs['activity_executor']._max_workers == 3
|
|
kwargs['activity_executor'].shutdown(wait=True, cancel_futures=True)
|
|
|
|
|
|
def test_prepare_worker_cleanup_queue_uses_cleanup_limits():
|
|
from model_manager.worker.prepare_worker import prepare_worker
|
|
from model_manager.workflows.cleanup_files import CleanupFiles
|
|
|
|
fake_worker = MagicMock()
|
|
fake_client = MagicMock()
|
|
fake_logger = MagicMock()
|
|
|
|
with patch(
|
|
'model_manager.worker.prepare_worker.Worker', return_value=fake_worker
|
|
) as worker_class:
|
|
with patch.dict(
|
|
'os.environ',
|
|
{
|
|
'CLEANUPFILES_ACTIVITY_EXECUTOR_MAX_WORKERS': '5',
|
|
'CLEANUPFILES_MAX_CONCURRENT_ACTIVITIES': '7',
|
|
},
|
|
clear=False,
|
|
):
|
|
worker = prepare_worker(
|
|
main_workflow=CleanupFiles,
|
|
other_workflows=[],
|
|
activities=[],
|
|
temporal_client=fake_client,
|
|
logger=fake_logger,
|
|
runtime='model-manager-worker',
|
|
)
|
|
|
|
assert worker is fake_worker
|
|
kwargs = worker_class.call_args.kwargs
|
|
assert kwargs['task_queue'] == 'cleanup_files-model-manager-worker-queue'
|
|
assert kwargs['max_concurrent_activities'] == 7
|
|
assert kwargs['activity_executor']._max_workers == 5
|
|
kwargs['activity_executor'].shutdown(wait=True, cancel_futures=True)
|