Code import - branch release/SIENTIAPDE-1645

2026-08-05 13:53:37 +00:00
commit d481e0acff
116 changed files with 92848 additions and 0 deletions
--- a/tests/worker/test_prepare_worker.py
+++ b/tests/worker/test_prepare_worker.py
@@ -0,0 +1,84 @@
+"""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)