0c9c5f300b
The worker imports sklearn (metrics, evidently) before PluginStore.install_runtime replaces it on disk with the runtime's pinned version, the same order that breaks every sklearn prediction in laborious. After installing the runtime, the worker now replaces its process once (SIENTIA_RUNTIME_READY=1 skips the second install), so every module is loaded from the runtime's packages.
1031 lines
37 KiB
Python
1031 lines
37 KiB
Python
"""Unit tests for worker module."""
|
|
|
|
import asyncio
|
|
import os
|
|
import sys
|
|
from unittest.mock import AsyncMock, Mock, patch
|
|
|
|
import pytest
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_env_vars():
|
|
"""Set up test environment variables."""
|
|
env_vars = {
|
|
'POD_ID': 'test-pod-123',
|
|
'HTTP_METRICS_PORT': '9090',
|
|
'HTTP_SDK_METRICS_PORT': '9091',
|
|
'TEMPORAL_HOST': 'localhost:7233',
|
|
'TEMPORAL_NAMESPACE': 'test-namespace',
|
|
'PROJECT_NAME': 'test-project',
|
|
'TRAIN_TASK_QUEUE': 'train_model-local_queue',
|
|
'CLEANUP_TASK_QUEUE': 'cleanup-local_queue',
|
|
'RUNTIME': 'model-manager-worker',
|
|
'STORE_BASE_URL': 'http://sientia-plugin-store.svc.cluster.local',
|
|
'STORE_OWNER': 'sientia',
|
|
'STORE_REPO': 'model-library-store',
|
|
}
|
|
|
|
with patch.dict(os.environ, env_vars, clear=False):
|
|
yield env_vars
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def runtime_ready(monkeypatch):
|
|
"""Boot as the re-executed process by default; the real os.execv must never run in tests."""
|
|
monkeypatch.setenv('SIENTIA_RUNTIME_READY', '1')
|
|
|
|
|
|
class _StopBootError(Exception):
|
|
"""Stops main() at a chosen point, e.g. where os.execv would replace the process."""
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_logger():
|
|
"""Create a mock logger."""
|
|
logger = Mock()
|
|
logger.custom_info = Mock()
|
|
logger.custom_error = Mock()
|
|
return logger
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_temporal_client():
|
|
"""Create a mock Temporal client."""
|
|
client_mock = AsyncMock()
|
|
client_mock.connect = AsyncMock()
|
|
return client_mock
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_worker():
|
|
"""Create a mock Temporal worker."""
|
|
worker_mock = Mock()
|
|
worker_mock.run = AsyncMock(return_value=None)
|
|
return worker_mock
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_notification_handler():
|
|
"""Create a mock notification handler."""
|
|
handler = Mock()
|
|
handler.shutdown = Mock()
|
|
return handler
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_activities():
|
|
"""Create a mock Activities instance."""
|
|
activities = AsyncMock()
|
|
activities.update_experiment_run = Mock()
|
|
activities.load_model_metadata = Mock()
|
|
activities.validate_train_params = Mock()
|
|
activities.train_model = Mock()
|
|
activities.cleanup_resources = Mock()
|
|
activities.shutdown = Mock()
|
|
return activities
|
|
|
|
|
|
def test_pod_id_from_env():
|
|
"""Test that POD_ID is correctly read from environment."""
|
|
with patch.dict(os.environ, {'POD_ID': 'pod-test-123'}):
|
|
# Re-import to get new env value
|
|
import importlib
|
|
|
|
import model_manager.worker.worker as worker_module
|
|
|
|
importlib.reload(worker_module)
|
|
|
|
assert worker_module.POD_ID == 'pod-test-123'
|
|
|
|
|
|
def test_sdk_metrics_port_default():
|
|
"""Test that SDK_METRICS_PORT uses default value."""
|
|
with patch.dict(os.environ, {}, clear=True):
|
|
import importlib
|
|
|
|
import model_manager.worker.worker as worker_module
|
|
|
|
importlib.reload(worker_module)
|
|
|
|
assert worker_module.SDK_METRICS_PORT == 9091
|
|
|
|
|
|
def test_sdk_metrics_port_from_env():
|
|
"""Test that SDK_METRICS_PORT is read from environment."""
|
|
with patch.dict(os.environ, {'HTTP_SDK_METRICS_PORT': '8888'}):
|
|
import importlib
|
|
|
|
import model_manager.worker.worker as worker_module
|
|
|
|
importlib.reload(worker_module)
|
|
|
|
assert worker_module.SDK_METRICS_PORT == 8888
|
|
|
|
|
|
@patch('model_manager.worker.worker.POD_ID', 'test-pod-123')
|
|
@patch('model_manager.worker.worker.start_http_server')
|
|
@patch('model_manager.worker.worker.metrics')
|
|
def test_start_prometheus_server_success(
|
|
mock_metrics, mock_start_http_server, mock_env_vars, mock_logger
|
|
):
|
|
"""Test successful Prometheus server startup."""
|
|
from model_manager.worker.worker import start_prometheus_server
|
|
|
|
mock_app_up = Mock()
|
|
mock_metrics.APP_UP.labels.return_value = mock_app_up
|
|
|
|
metadata: dict[str, str | None] = {'pod_id': 'test-pod-123', 'workflow_name': 'train_model'}
|
|
|
|
start_prometheus_server(mock_logger, metadata)
|
|
|
|
# Verify HTTP server started
|
|
mock_start_http_server.assert_called_once_with(9090)
|
|
|
|
# Verify APP_UP metric was set to 1
|
|
mock_metrics.APP_UP.labels.assert_called_once_with(pod_id='test-pod-123')
|
|
mock_app_up.set.assert_called_once_with(1)
|
|
mock_logger.custom_info.assert_called_once()
|
|
|
|
|
|
@patch('model_manager.worker.worker.start_http_server')
|
|
@patch('model_manager.worker.worker.metrics')
|
|
def test_start_prometheus_server_custom_port(mock_metrics, mock_start_http_server, mock_logger):
|
|
"""Test Prometheus server startup with custom port."""
|
|
from model_manager.worker.worker import start_prometheus_server
|
|
|
|
with patch.dict(os.environ, {'HTTP_METRICS_PORT': '8080', 'POD_ID': 'custom-pod'}):
|
|
mock_app_up = Mock()
|
|
mock_metrics.APP_UP.labels.return_value = mock_app_up
|
|
|
|
metadata: dict[str, str | None] = {'pod_id': 'custom-pod', 'workflow_name': 'train_model'}
|
|
|
|
start_prometheus_server(mock_logger, metadata)
|
|
|
|
mock_start_http_server.assert_called_once_with(8080)
|
|
|
|
|
|
@patch('model_manager.worker.worker.start_http_server')
|
|
@patch('model_manager.worker.worker.metrics')
|
|
@patch('model_manager.worker.worker.os._exit')
|
|
def test_start_prometheus_server_failure(
|
|
mock_exit, mock_metrics, mock_start_http_server, mock_env_vars, mock_logger
|
|
):
|
|
"""Test Prometheus server startup failure."""
|
|
from model_manager.worker.worker import start_prometheus_server
|
|
|
|
mock_start_http_server.side_effect = OSError('Port already in use')
|
|
|
|
metadata: dict[str, str | None] = {'pod_id': 'test-pod-123', 'workflow_name': 'train_model'}
|
|
|
|
start_prometheus_server(mock_logger, metadata)
|
|
|
|
# Verify exit was called with code 1 e log crítico emitido
|
|
mock_exit.assert_called_once_with(1)
|
|
mock_logger.custom_critical.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch('model_manager.worker.worker.RUNTIME', 'model-manager-worker')
|
|
@patch('model_manager.worker.worker.prepare_worker')
|
|
@patch('model_manager.worker.worker.client.Client')
|
|
@patch('model_manager.worker.worker.Runtime')
|
|
@patch('model_manager.worker.worker.Activities')
|
|
@patch('model_manager.worker.worker.NotificationHandler')
|
|
@patch('model_manager.worker.worker.PluginStore')
|
|
@patch('model_manager.worker.worker.build_plugin_store_config')
|
|
@patch('model_manager.worker.worker.build_mongodb_config')
|
|
@patch('model_manager.worker.worker.build_postgres_config')
|
|
@patch('model_manager.worker.worker.build_mlflow_config')
|
|
@patch('model_manager.worker.worker.build_minio_config')
|
|
@patch('model_manager.worker.worker.get_logger')
|
|
@patch('model_manager.worker.worker.start_prometheus_server')
|
|
@patch('model_manager.worker.worker.metrics')
|
|
@patch('model_manager.worker.worker.ensure_runtime_directories')
|
|
async def test_main_successful_startup(
|
|
mock_ensure_runtime_directories,
|
|
mock_metrics,
|
|
mock_start_prometheus,
|
|
mock_get_logger,
|
|
mock_build_minio,
|
|
mock_build_mlflow,
|
|
mock_build_postgres,
|
|
mock_build_mongodb,
|
|
mock_build_plugin_store_config,
|
|
mock_plugin_store_class,
|
|
mock_notification_handler_class,
|
|
mock_activities_class,
|
|
mock_runtime_class,
|
|
mock_client_class,
|
|
mock_prepare_worker,
|
|
mock_env_vars,
|
|
mock_logger,
|
|
mock_temporal_client,
|
|
mock_worker,
|
|
mock_notification_handler,
|
|
mock_activities,
|
|
):
|
|
"""Test successful main() execution until workers start."""
|
|
from model_manager.worker.worker import main
|
|
|
|
# Setup mocks
|
|
mock_get_logger.return_value = mock_logger
|
|
mock_build_mongodb.return_value = {
|
|
'connection_string': 'mongodb://test',
|
|
'database_name': 'test_db',
|
|
'uri': 'localhost:27018',
|
|
}
|
|
mock_build_postgres.return_value = {}
|
|
mock_build_mlflow.return_value = {}
|
|
mock_build_minio.return_value = {}
|
|
|
|
mock_notification_handler_class.return_value = mock_notification_handler
|
|
mock_activities_class.return_value = mock_activities
|
|
|
|
mock_plugin_store_instance = AsyncMock()
|
|
mock_plugin_store_instance.install_runtime = AsyncMock(
|
|
return_value={'runtime': 'model-manager-worker', 'installed': []},
|
|
)
|
|
mock_plugin_store_class.return_value = mock_plugin_store_instance
|
|
mock_build_plugin_store_config.return_value = {
|
|
'base_url': 'http://sientia-plugin-store.svc.cluster.local',
|
|
'owner': 'sientia',
|
|
'repo': 'model-library-store',
|
|
'branch': 'main',
|
|
'username': 'gitea-user',
|
|
'password': 'gitea-password',
|
|
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
|
|
'pypi_username': None,
|
|
'pypi_password': None,
|
|
'cache_ttl_seconds': None,
|
|
}
|
|
|
|
mock_runtime = Mock()
|
|
mock_runtime_class.return_value = mock_runtime
|
|
|
|
mock_client_instance = AsyncMock()
|
|
mock_client_instance.config = Mock(return_value={'plugins': [], 'interceptors': []})
|
|
mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
|
|
|
|
mock_worker_instance = Mock()
|
|
mock_worker_instance.run = AsyncMock(
|
|
side_effect=asyncio.CancelledError()
|
|
) # Simulate interruption
|
|
mock_prepare_worker.return_value = mock_worker_instance
|
|
|
|
mock_app_up = Mock()
|
|
mock_metrics.APP_UP.labels.return_value = mock_app_up
|
|
|
|
# Run main() and expect it to exit due to CancelledError
|
|
with pytest.raises(SystemExit) as exc_info:
|
|
await main()
|
|
|
|
assert exc_info.value.code == 1
|
|
|
|
# Verify all initialization steps were called
|
|
mock_get_logger.assert_called_once()
|
|
mock_start_prometheus.assert_called_once()
|
|
mock_notification_handler_class.assert_called_once()
|
|
mock_activities_class.assert_called_once()
|
|
mock_client_class.connect.assert_called_once()
|
|
# Agora são criados dois Workers: um para train_model-queue e outro para cleanup-queue
|
|
assert mock_prepare_worker.call_count == 2
|
|
|
|
# Verify cleanup was performed
|
|
mock_notification_handler.shutdown.assert_called_once()
|
|
mock_activities.shutdown.assert_called_once()
|
|
mock_app_up.set.assert_called_with(0)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch('model_manager.worker.worker.RUNTIME', 'model-manager-worker')
|
|
@patch('model_manager.worker.worker.prepare_worker')
|
|
@patch('model_manager.worker.worker.client.Client')
|
|
@patch('model_manager.worker.worker.Runtime')
|
|
@patch('model_manager.worker.worker.Activities')
|
|
@patch('model_manager.worker.worker.NotificationHandler')
|
|
@patch('model_manager.worker.worker.PluginStore')
|
|
@patch('model_manager.worker.worker.build_plugin_store_config')
|
|
@patch('model_manager.worker.worker.build_mongodb_config')
|
|
@patch('model_manager.worker.worker.build_postgres_config')
|
|
@patch('model_manager.worker.worker.build_mlflow_config')
|
|
@patch('model_manager.worker.worker.build_minio_config')
|
|
@patch('model_manager.worker.worker.get_logger')
|
|
@patch('model_manager.worker.worker.start_prometheus_server')
|
|
@patch('model_manager.worker.worker.metrics')
|
|
@patch('model_manager.worker.worker.ensure_runtime_directories')
|
|
async def test_main_handles_exception(
|
|
mock_ensure_runtime_directories,
|
|
mock_metrics,
|
|
mock_start_prometheus,
|
|
mock_get_logger,
|
|
mock_build_minio,
|
|
mock_build_mlflow,
|
|
mock_build_postgres,
|
|
mock_build_mongodb,
|
|
mock_build_plugin_store_config,
|
|
mock_plugin_store_class,
|
|
mock_notification_handler_class,
|
|
mock_activities_class,
|
|
mock_runtime_class,
|
|
mock_client_class,
|
|
mock_prepare_worker,
|
|
mock_env_vars,
|
|
mock_logger,
|
|
):
|
|
"""Test main() handles exceptions and performs cleanup."""
|
|
from model_manager.worker.worker import main
|
|
|
|
# Setup mocks
|
|
mock_get_logger.return_value = mock_logger
|
|
mock_build_mongodb.return_value = {
|
|
'connection_string': 'mongodb://test',
|
|
'database_name': 'test_db',
|
|
'uri': 'localhost:27018',
|
|
}
|
|
mock_build_postgres.return_value = {}
|
|
mock_build_mlflow.return_value = {}
|
|
mock_build_minio.return_value = {}
|
|
|
|
mock_notification_handler = Mock()
|
|
mock_notification_handler.shutdown = Mock()
|
|
mock_notification_handler_class.return_value = mock_notification_handler
|
|
|
|
mock_activities = AsyncMock()
|
|
mock_activities.shutdown = Mock()
|
|
mock_activities_class.return_value = mock_activities
|
|
|
|
mock_runtime = Mock()
|
|
mock_runtime_class.return_value = mock_runtime
|
|
|
|
mock_client_instance = AsyncMock()
|
|
mock_client_instance.config = Mock(return_value={'plugins': [], 'interceptors': []})
|
|
mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
|
|
|
|
mock_worker_instance = Mock()
|
|
mock_worker_instance.run = AsyncMock(side_effect=RuntimeError('Worker failed'))
|
|
mock_prepare_worker.return_value = mock_worker_instance
|
|
|
|
mock_app_up = Mock()
|
|
mock_metrics.APP_UP.labels.return_value = mock_app_up
|
|
|
|
mock_plugin_store_instance = AsyncMock()
|
|
mock_plugin_store_instance.install_runtime = AsyncMock(
|
|
return_value={'runtime': 'model-manager-worker', 'installed': []},
|
|
)
|
|
mock_plugin_store_class.return_value = mock_plugin_store_instance
|
|
mock_build_plugin_store_config.return_value = {
|
|
'base_url': 'http://sientia-plugin-store.svc.cluster.local',
|
|
'owner': 'sientia',
|
|
'repo': 'model-library-store',
|
|
'branch': 'main',
|
|
'username': 'gitea-user',
|
|
'password': 'gitea-password',
|
|
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
|
|
'pypi_username': None,
|
|
'pypi_password': None,
|
|
'cache_ttl_seconds': None,
|
|
}
|
|
|
|
# Run main() and expect SystemExit
|
|
with pytest.raises(SystemExit) as exc_info:
|
|
await main()
|
|
|
|
assert exc_info.value.code == 1
|
|
|
|
# Verify error was logged
|
|
mock_logger.custom_error.assert_called_once()
|
|
assert 'Worker failed' in str(mock_logger.custom_error.call_args)
|
|
|
|
# Verify cleanup was performed
|
|
mock_notification_handler.shutdown.assert_called_once()
|
|
mock_activities.shutdown.assert_called_once()
|
|
mock_app_up.set.assert_called_with(0)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch('model_manager.worker.worker.RUNTIME', 'model-manager-worker')
|
|
@patch('model_manager.worker.worker.create_cleanup_schedule')
|
|
@patch('model_manager.worker.worker.prepare_worker')
|
|
@patch('model_manager.worker.worker.client.Client')
|
|
@patch('model_manager.worker.worker.Runtime')
|
|
@patch('model_manager.worker.worker.Activities')
|
|
@patch('model_manager.worker.worker.NotificationHandler')
|
|
@patch('model_manager.worker.worker.PluginStore')
|
|
@patch('model_manager.worker.worker.build_plugin_store_config')
|
|
@patch('model_manager.worker.worker.build_mongodb_config')
|
|
@patch('model_manager.worker.worker.build_postgres_config')
|
|
@patch('model_manager.worker.worker.build_mlflow_config')
|
|
@patch('model_manager.worker.worker.build_minio_config')
|
|
@patch('model_manager.worker.worker.get_logger')
|
|
@patch('model_manager.worker.worker.start_prometheus_server')
|
|
@patch('model_manager.worker.worker.metrics')
|
|
@patch('model_manager.worker.worker.ensure_runtime_directories')
|
|
async def test_main_temporal_client_configuration(
|
|
mock_ensure_runtime_directories,
|
|
mock_metrics,
|
|
mock_start_prometheus,
|
|
mock_get_logger,
|
|
mock_build_minio,
|
|
mock_build_mlflow,
|
|
mock_build_postgres,
|
|
mock_build_mongodb,
|
|
mock_build_plugin_store_config,
|
|
mock_plugin_store_class,
|
|
mock_notification_handler_class,
|
|
mock_activities_class,
|
|
mock_runtime_class,
|
|
mock_client_class,
|
|
mock_prepare_worker,
|
|
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',
|
|
'TEMPORAL_USE_TLS': 'true',
|
|
'RUNTIME': 'model-manager-worker',
|
|
'STORE_BASE_URL': 'http://sientia-plugin-store.svc.cluster.local',
|
|
'STORE_OWNER': 'sientia',
|
|
'STORE_REPO': 'model-library-store',
|
|
},
|
|
):
|
|
# Setup mocks
|
|
mock_get_logger.return_value = mock_logger
|
|
mock_build_mongodb.return_value = {
|
|
'connection_string': 'mongodb://test',
|
|
'database_name': 'test_db',
|
|
'uri': 'localhost:27018',
|
|
}
|
|
mock_build_postgres.return_value = {}
|
|
mock_build_mlflow.return_value = {}
|
|
mock_build_minio.return_value = {}
|
|
|
|
mock_plugin_store_instance = AsyncMock()
|
|
mock_plugin_store_instance.install_runtime = AsyncMock(
|
|
return_value={'runtime': 'model-manager-worker', 'installed': []},
|
|
)
|
|
mock_plugin_store_class.return_value = mock_plugin_store_instance
|
|
mock_build_plugin_store_config.return_value = {
|
|
'base_url': 'http://sientia-plugin-store.svc.cluster.local',
|
|
'owner': 'sientia',
|
|
'repo': 'model-library-store',
|
|
'branch': 'main',
|
|
'username': 'gitea-user',
|
|
'password': 'gitea-password',
|
|
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
|
|
'pypi_username': None,
|
|
'pypi_password': None,
|
|
'cache_ttl_seconds': None,
|
|
}
|
|
|
|
mock_notification_handler = Mock()
|
|
mock_notification_handler.shutdown = Mock()
|
|
mock_notification_handler_class.return_value = mock_notification_handler
|
|
|
|
mock_activities = AsyncMock()
|
|
mock_activities.shutdown = Mock()
|
|
mock_activities_class.return_value = mock_activities
|
|
|
|
mock_runtime = Mock()
|
|
mock_runtime_class.return_value = mock_runtime
|
|
|
|
mock_client_instance = AsyncMock()
|
|
mock_client_instance.config = Mock(return_value={'plugins': [], 'interceptors': []})
|
|
mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
|
|
|
|
mock_worker_instance = Mock()
|
|
mock_worker_instance.run = AsyncMock(side_effect=asyncio.CancelledError())
|
|
mock_prepare_worker.return_value = mock_worker_instance
|
|
|
|
mock_app_up = Mock()
|
|
mock_metrics.APP_UP.labels.return_value = mock_app_up
|
|
|
|
# Run main()
|
|
with pytest.raises(SystemExit):
|
|
await main()
|
|
|
|
# Verify Temporal client was configured with correct parameters
|
|
mock_client_class.connect.assert_called_once_with(
|
|
target_host='temporal.example.com:7233',
|
|
namespace='production',
|
|
runtime=mock_runtime,
|
|
tls=True,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch('model_manager.worker.worker.RUNTIME', 'model-manager-worker')
|
|
@patch('model_manager.worker.worker.create_cleanup_schedule')
|
|
@patch('model_manager.worker.worker.prepare_worker')
|
|
@patch('model_manager.worker.worker.client.Client')
|
|
@patch('model_manager.worker.worker.Runtime')
|
|
@patch('model_manager.worker.worker.Activities')
|
|
@patch('model_manager.worker.worker.NotificationHandler')
|
|
@patch('model_manager.worker.worker.PluginStore')
|
|
@patch('model_manager.worker.worker.build_plugin_store_config')
|
|
@patch('model_manager.worker.worker.build_mongodb_config')
|
|
@patch('model_manager.worker.worker.build_postgres_config')
|
|
@patch('model_manager.worker.worker.build_mlflow_config')
|
|
@patch('model_manager.worker.worker.build_minio_config')
|
|
@patch('model_manager.worker.worker.get_logger')
|
|
@patch('model_manager.worker.worker.start_prometheus_server')
|
|
@patch('model_manager.worker.worker.metrics')
|
|
@patch('model_manager.worker.worker.ensure_runtime_directories')
|
|
async def test_main_worker_configuration(
|
|
mock_ensure_runtime_directories,
|
|
mock_metrics,
|
|
mock_start_prometheus,
|
|
mock_get_logger,
|
|
mock_build_minio,
|
|
mock_build_mlflow,
|
|
mock_build_postgres,
|
|
mock_build_mongodb,
|
|
mock_build_plugin_store_config,
|
|
mock_plugin_store_class,
|
|
mock_notification_handler_class,
|
|
mock_activities_class,
|
|
mock_runtime_class,
|
|
mock_client_class,
|
|
mock_prepare_worker,
|
|
mock_create_cleanup_schedule,
|
|
mock_env_vars,
|
|
mock_logger,
|
|
):
|
|
"""Test that prepare_worker is configured with correct workflows and activities."""
|
|
from model_manager.worker.worker import main
|
|
from model_manager.workflows.cleanup_files import CleanupFiles
|
|
from model_manager.workflows.train_model import TrainModel
|
|
|
|
mock_create_cleanup_schedule.return_value = AsyncMock()
|
|
|
|
mock_get_logger.return_value = mock_logger
|
|
mock_build_mongodb.return_value = {
|
|
'connection_string': 'mongodb://test',
|
|
'database_name': 'test_db',
|
|
'uri': 'localhost:27018',
|
|
}
|
|
mock_build_postgres.return_value = {}
|
|
mock_build_mlflow.return_value = {}
|
|
mock_build_minio.return_value = {}
|
|
|
|
mock_notification_handler = Mock()
|
|
mock_notification_handler_class.return_value = mock_notification_handler
|
|
|
|
mock_activities = AsyncMock()
|
|
mock_activities.update_experiment_run = Mock()
|
|
mock_activities.load_model_metadata = Mock()
|
|
mock_activities.validate_train_params = Mock()
|
|
mock_activities.train_model = Mock()
|
|
mock_activities.cleanup_resources = Mock()
|
|
mock_activities.cleanup_temp_directories = Mock()
|
|
mock_activities.shutdown = Mock()
|
|
mock_activities_class.return_value = mock_activities
|
|
|
|
mock_runtime = Mock()
|
|
mock_runtime_class.return_value = mock_runtime
|
|
|
|
mock_client_instance = AsyncMock()
|
|
mock_client_instance.config = Mock(return_value={'plugins': [], 'interceptors': []})
|
|
mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
|
|
|
|
mock_worker_instance = Mock()
|
|
mock_worker_instance.run = AsyncMock(side_effect=asyncio.CancelledError())
|
|
mock_prepare_worker.return_value = mock_worker_instance
|
|
|
|
mock_app_up = Mock()
|
|
mock_metrics.APP_UP.labels.return_value = mock_app_up
|
|
|
|
mock_plugin_store_instance = AsyncMock()
|
|
mock_plugin_store_instance.install_runtime = AsyncMock(
|
|
return_value={'runtime': 'model-manager-worker', 'installed': []},
|
|
)
|
|
mock_plugin_store_class.return_value = mock_plugin_store_instance
|
|
mock_build_plugin_store_config.return_value = {
|
|
'base_url': 'http://sientia-plugin-store.svc.cluster.local',
|
|
'owner': 'sientia',
|
|
'repo': 'model-library-store',
|
|
'branch': 'main',
|
|
'username': 'gitea-user',
|
|
'password': 'gitea-password',
|
|
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
|
|
'pypi_username': None,
|
|
'pypi_password': None,
|
|
'cache_ttl_seconds': None,
|
|
}
|
|
|
|
with pytest.raises(SystemExit):
|
|
await main()
|
|
|
|
assert mock_prepare_worker.call_count == 2
|
|
|
|
train_call = mock_prepare_worker.call_args_list[0]
|
|
assert train_call.kwargs['temporal_client'] is mock_client_instance
|
|
assert train_call.kwargs['logger'] is mock_logger
|
|
assert train_call.kwargs['main_workflow'] is TrainModel
|
|
assert train_call.kwargs['other_workflows'] == []
|
|
train_activities_list = train_call.kwargs['activities']
|
|
assert mock_activities.update_experiment_run in train_activities_list
|
|
assert mock_activities.load_model_metadata in train_activities_list
|
|
assert mock_activities.validate_train_params in train_activities_list
|
|
assert mock_activities.train_model in train_activities_list
|
|
assert mock_activities.cleanup_resources in train_activities_list
|
|
|
|
cleanup_call = mock_prepare_worker.call_args_list[1]
|
|
assert cleanup_call.kwargs['temporal_client'] is mock_client_instance
|
|
assert cleanup_call.kwargs['logger'] is mock_logger
|
|
assert cleanup_call.kwargs['main_workflow'] is CleanupFiles
|
|
assert cleanup_call.kwargs['other_workflows'] == []
|
|
assert cleanup_call.kwargs['activities'] == [mock_activities.cleanup_temp_directories]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch('model_manager.worker.worker.RUNTIME', 'model-manager-worker')
|
|
@patch('model_manager.worker.worker.create_cleanup_schedule')
|
|
@patch('model_manager.worker.worker.prepare_worker')
|
|
@patch('model_manager.worker.worker.client.Client')
|
|
@patch('model_manager.worker.worker.Runtime')
|
|
@patch('model_manager.worker.worker.Activities')
|
|
@patch('model_manager.worker.worker.NotificationHandler')
|
|
@patch('model_manager.worker.worker.PluginStore')
|
|
@patch('model_manager.worker.worker.build_plugin_store_config')
|
|
@patch('model_manager.worker.worker.build_mongodb_config')
|
|
@patch('model_manager.worker.worker.build_postgres_config')
|
|
@patch('model_manager.worker.worker.build_mlflow_config')
|
|
@patch('model_manager.worker.worker.build_minio_config')
|
|
@patch('model_manager.worker.worker.get_logger')
|
|
@patch('model_manager.worker.worker.start_prometheus_server')
|
|
@patch('model_manager.worker.worker.metrics')
|
|
@patch('model_manager.worker.worker.ensure_runtime_directories')
|
|
async def test_main_schedule_creation_failure_does_not_stop_worker(
|
|
mock_ensure_runtime_directories,
|
|
mock_metrics,
|
|
mock_start_prometheus,
|
|
mock_get_logger,
|
|
mock_build_minio,
|
|
mock_build_mlflow,
|
|
mock_build_postgres,
|
|
mock_build_mongodb,
|
|
mock_build_plugin_store_config,
|
|
mock_plugin_store_class,
|
|
mock_notification_handler_class,
|
|
mock_activities_class,
|
|
mock_runtime_class,
|
|
mock_client_class,
|
|
mock_prepare_worker,
|
|
mock_create_cleanup_schedule,
|
|
mock_logger,
|
|
):
|
|
"""Test that schedule creation failure does not prevent worker startup."""
|
|
from model_manager.worker.worker import main
|
|
|
|
# Mock schedule creation to raise an exception (as coroutine)
|
|
async def mock_schedule_error(*args, **kwargs):
|
|
raise Exception('Schedule creation failed')
|
|
|
|
mock_create_cleanup_schedule.side_effect = mock_schedule_error
|
|
|
|
# Setup mocks
|
|
mock_get_logger.return_value = mock_logger
|
|
mock_build_mongodb.return_value = {
|
|
'connection_string': 'mongodb://test',
|
|
'database_name': 'test_db',
|
|
'uri': 'localhost:27018',
|
|
}
|
|
mock_build_postgres.return_value = {}
|
|
mock_build_mlflow.return_value = {}
|
|
mock_build_minio.return_value = {}
|
|
|
|
mock_notification_handler = Mock()
|
|
mock_notification_handler.shutdown = Mock()
|
|
mock_notification_handler_class.return_value = mock_notification_handler
|
|
|
|
mock_activities = AsyncMock()
|
|
mock_activities.shutdown = Mock()
|
|
mock_activities_class.return_value = mock_activities
|
|
|
|
mock_runtime = Mock()
|
|
mock_runtime_class.return_value = mock_runtime
|
|
|
|
mock_client_instance = AsyncMock()
|
|
mock_client_instance.config = Mock(return_value={'plugins': [], 'interceptors': []})
|
|
mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
|
|
|
|
mock_worker_instance = Mock()
|
|
mock_worker_instance.run = AsyncMock(side_effect=asyncio.CancelledError())
|
|
mock_prepare_worker.return_value = mock_worker_instance
|
|
|
|
mock_app_up = Mock()
|
|
mock_metrics.APP_UP.labels.return_value = mock_app_up
|
|
|
|
mock_plugin_store_instance = AsyncMock()
|
|
mock_plugin_store_instance.install_runtime = AsyncMock(
|
|
return_value={'runtime': 'model-manager-worker', 'installed': []},
|
|
)
|
|
mock_plugin_store_class.return_value = mock_plugin_store_instance
|
|
mock_build_plugin_store_config.return_value = {
|
|
'base_url': 'http://sientia-plugin-store.svc.cluster.local',
|
|
'owner': 'sientia',
|
|
'repo': 'model-library-store',
|
|
'branch': 'main',
|
|
'username': 'gitea-user',
|
|
'password': 'gitea-password',
|
|
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
|
|
'pypi_username': None,
|
|
'pypi_password': None,
|
|
'cache_ttl_seconds': None,
|
|
}
|
|
|
|
with pytest.raises(SystemExit):
|
|
await main()
|
|
|
|
mock_create_cleanup_schedule.assert_called_once()
|
|
|
|
schedule_error_logged = False
|
|
for call in mock_logger.custom_error.call_args_list:
|
|
if call[0] and 'Failed to configure cleanup schedule' in call[0][0]:
|
|
schedule_error_logged = True
|
|
break
|
|
assert schedule_error_logged, 'Schedule creation error should be logged'
|
|
|
|
assert mock_prepare_worker.call_count == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch('model_manager.worker.worker.RUNTIME', None)
|
|
@patch('model_manager.worker.worker.prepare_worker')
|
|
@patch('model_manager.worker.worker.client.Client')
|
|
@patch('model_manager.worker.worker.Runtime')
|
|
@patch('model_manager.worker.worker.Activities')
|
|
@patch('model_manager.worker.worker.NotificationHandler')
|
|
@patch('model_manager.worker.worker.PluginStore')
|
|
@patch('model_manager.worker.worker.build_plugin_store_config')
|
|
@patch('model_manager.worker.worker.build_mongodb_config')
|
|
@patch('model_manager.worker.worker.build_postgres_config')
|
|
@patch('model_manager.worker.worker.build_mlflow_config')
|
|
@patch('model_manager.worker.worker.build_minio_config')
|
|
@patch('model_manager.worker.worker.get_logger')
|
|
@patch('model_manager.worker.worker.start_prometheus_server')
|
|
@patch('model_manager.worker.worker.metrics')
|
|
@patch('model_manager.worker.worker.ensure_runtime_directories')
|
|
async def test_main_missing_runtime_uses_single_fallback(
|
|
mock_ensure_runtime_directories,
|
|
mock_metrics,
|
|
mock_start_prometheus,
|
|
mock_get_logger,
|
|
mock_build_minio,
|
|
mock_build_mlflow,
|
|
mock_build_postgres,
|
|
mock_build_mongodb,
|
|
mock_build_plugin_store_config,
|
|
mock_plugin_store_class,
|
|
mock_notification_handler_class,
|
|
mock_activities_class,
|
|
mock_runtime_class,
|
|
mock_client_class,
|
|
mock_prepare_worker,
|
|
mock_logger,
|
|
):
|
|
"""Test that main() uses single runtime fallback when RUNTIME is missing."""
|
|
from model_manager.worker.worker import main
|
|
|
|
mock_get_logger.return_value = mock_logger
|
|
mock_build_mongodb.return_value = {
|
|
'connection_string': 'mongodb://test',
|
|
'database_name': 'test_db',
|
|
'uri': 'localhost:27018',
|
|
}
|
|
mock_build_postgres.return_value = {}
|
|
mock_build_mlflow.return_value = {}
|
|
mock_build_minio.return_value = {}
|
|
|
|
mock_notification_handler = Mock()
|
|
mock_notification_handler.shutdown = Mock()
|
|
mock_notification_handler_class.return_value = mock_notification_handler
|
|
|
|
mock_activities = AsyncMock()
|
|
mock_activities.shutdown = Mock()
|
|
mock_activities_class.return_value = mock_activities
|
|
|
|
mock_runtime = Mock()
|
|
mock_runtime_class.return_value = mock_runtime
|
|
|
|
mock_client_instance = AsyncMock()
|
|
mock_client_instance.config = Mock(return_value={'plugins': [], 'interceptors': []})
|
|
mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
|
|
|
|
mock_worker_instance = Mock()
|
|
mock_worker_instance.run = AsyncMock(side_effect=asyncio.CancelledError())
|
|
mock_prepare_worker.return_value = mock_worker_instance
|
|
|
|
mock_plugin_store_instance = AsyncMock()
|
|
mock_plugin_store_instance.install_runtime = AsyncMock(
|
|
return_value={'runtime': 'single', 'installed': []},
|
|
)
|
|
mock_plugin_store_class.return_value = mock_plugin_store_instance
|
|
mock_build_plugin_store_config.return_value = {
|
|
'base_url': 'http://sientia-plugin-store.svc.cluster.local',
|
|
'owner': 'sientia',
|
|
'repo': 'model-library-store',
|
|
'branch': 'main',
|
|
'username': 'gitea-user',
|
|
'password': 'gitea-password',
|
|
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
|
|
'pypi_username': None,
|
|
'pypi_password': None,
|
|
'cache_ttl_seconds': None,
|
|
}
|
|
|
|
mock_app_up = Mock()
|
|
mock_metrics.APP_UP.labels.return_value = mock_app_up
|
|
|
|
with pytest.raises(SystemExit):
|
|
await main()
|
|
|
|
assert mock_prepare_worker.call_count == 2
|
|
assert mock_prepare_worker.call_args_list[0].kwargs['runtime'] == 'single'
|
|
assert mock_prepare_worker.call_args_list[1].kwargs['runtime'] == 'single'
|
|
|
|
|
|
_PLUGIN_STORE_CONFIG = {
|
|
'base_url': 'http://sientia-plugin-store.svc.cluster.local',
|
|
'owner': 'sientia',
|
|
'repo': 'model-library-store',
|
|
'branch': 'main',
|
|
'username': 'gitea-user',
|
|
'password': 'gitea-password',
|
|
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
|
|
'pypi_username': None,
|
|
'pypi_password': None,
|
|
'cache_ttl_seconds': None,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch('model_manager.worker.worker.RUNTIME', 'basic')
|
|
# Stops boot at Activities, so the red phase does not try to reach a real Temporal server.
|
|
@patch('model_manager.worker.worker.Activities', side_effect=_StopBootError)
|
|
@patch('model_manager.worker.worker.NotificationHandler')
|
|
@patch('model_manager.worker.worker.PluginStore')
|
|
@patch('model_manager.worker.worker.build_plugin_store_config', return_value=_PLUGIN_STORE_CONFIG)
|
|
@patch('model_manager.worker.worker.build_postgres_config', return_value={})
|
|
@patch('model_manager.worker.worker.build_mlflow_config', return_value={})
|
|
@patch('model_manager.worker.worker.build_minio_config', return_value={})
|
|
@patch(
|
|
'model_manager.worker.worker.build_mongodb_config',
|
|
return_value={'connection_string': 'mongodb://test', 'database_name': 'db', 'uri': 'u'},
|
|
)
|
|
@patch('model_manager.worker.worker.get_logger')
|
|
@patch('model_manager.worker.worker.start_prometheus_server')
|
|
@patch('model_manager.worker.worker.ensure_runtime_directories')
|
|
async def test_main_reexecs_after_installing_runtime(
|
|
_mock_dirs,
|
|
_mock_prometheus,
|
|
_mock_get_logger,
|
|
_mock_mongo,
|
|
_mock_minio,
|
|
_mock_mlflow,
|
|
_mock_postgres,
|
|
_mock_ps_config,
|
|
mock_plugin_store_class,
|
|
_mock_notification,
|
|
mock_activities_class,
|
|
monkeypatch,
|
|
):
|
|
"""First boot installs the runtime, then replaces the process so imports see it on disk."""
|
|
from model_manager.worker import worker as worker_module
|
|
|
|
monkeypatch.setenv('SIENTIA_RUNTIME_READY', '0')
|
|
mock_plugin_store_class.return_value.install_runtime = AsyncMock(return_value=None)
|
|
|
|
with (
|
|
patch('model_manager.worker.worker.os.execv', side_effect=_StopBootError) as m_execv,
|
|
patch.object(sys, 'orig_argv', ['python', '-m', 'model_manager.worker.worker']),
|
|
pytest.raises(_StopBootError),
|
|
):
|
|
await worker_module.main()
|
|
|
|
mock_plugin_store_class.return_value.install_runtime.assert_awaited_once_with(
|
|
runtime_name='basic'
|
|
)
|
|
assert os.environ['SIENTIA_RUNTIME_READY'] == '1'
|
|
m_execv.assert_called_once_with(
|
|
sys.executable, [sys.executable, '-m', 'model_manager.worker.worker']
|
|
)
|
|
mock_activities_class.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch('model_manager.worker.worker.RUNTIME', 'basic')
|
|
@patch('model_manager.worker.worker.Activities', side_effect=_StopBootError)
|
|
@patch('model_manager.worker.worker.NotificationHandler')
|
|
@patch('model_manager.worker.worker.PluginStore')
|
|
@patch('model_manager.worker.worker.build_plugin_store_config', return_value=_PLUGIN_STORE_CONFIG)
|
|
@patch('model_manager.worker.worker.build_postgres_config', return_value={})
|
|
@patch('model_manager.worker.worker.build_mlflow_config', return_value={})
|
|
@patch('model_manager.worker.worker.build_minio_config', return_value={})
|
|
@patch(
|
|
'model_manager.worker.worker.build_mongodb_config',
|
|
return_value={'connection_string': 'mongodb://test', 'database_name': 'db', 'uri': 'u'},
|
|
)
|
|
@patch('model_manager.worker.worker.get_logger')
|
|
@patch('model_manager.worker.worker.start_prometheus_server')
|
|
@patch('model_manager.worker.worker.ensure_runtime_directories')
|
|
async def test_main_skips_install_when_runtime_ready(
|
|
_mock_dirs,
|
|
_mock_prometheus,
|
|
_mock_get_logger,
|
|
_mock_mongo,
|
|
_mock_minio,
|
|
_mock_mlflow,
|
|
_mock_postgres,
|
|
_mock_ps_config,
|
|
mock_plugin_store_class,
|
|
_mock_notification,
|
|
mock_activities_class,
|
|
):
|
|
"""The re-executed process skips the install and the exec and goes on booting."""
|
|
from model_manager.worker import worker as worker_module
|
|
|
|
mock_plugin_store_class.return_value.install_runtime = AsyncMock(return_value=None)
|
|
|
|
with (
|
|
patch('model_manager.worker.worker.os.execv') as m_execv,
|
|
pytest.raises(_StopBootError),
|
|
):
|
|
await worker_module.main()
|
|
|
|
mock_plugin_store_class.return_value.install_runtime.assert_not_awaited()
|
|
m_execv.assert_not_called()
|
|
mock_activities_class.assert_called_once()
|
|
|
|
|
|
@patch('model_manager.worker.worker.asyncio.run')
|
|
def test_main_entrypoint(mock_asyncio_run):
|
|
"""Test the __main__ entrypoint."""
|
|
# Import and execute the main block
|
|
with patch.object(sys, 'argv', ['worker.py']):
|
|
import model_manager.worker.worker as worker_module
|
|
|
|
# Simulate running the module
|
|
worker_module.main = AsyncMock()
|
|
|
|
# This would normally be called by asyncio.run(main())
|
|
# We just verify the pattern is correct
|
|
assert callable(worker_module.main)
|
|
|
|
|
|
def test_worker_module_docstring():
|
|
"""Test that worker module has comprehensive documentation."""
|
|
import model_manager.worker.worker as worker_module
|
|
|
|
assert worker_module.__doc__ is not None
|
|
assert 'Temporal' in worker_module.__doc__
|
|
assert 'worker' in worker_module.__doc__
|
|
|
|
|
|
@patch('model_manager.worker.worker.start_http_server')
|
|
@patch('model_manager.worker.worker.metrics')
|
|
def test_start_prometheus_server_prints_success(
|
|
mock_metrics, mock_start_http_server, capsys, mock_env_vars, mock_logger
|
|
):
|
|
"""Test that start_prometheus_server prints success message."""
|
|
from model_manager.worker.worker import start_prometheus_server
|
|
|
|
mock_app_up = Mock()
|
|
mock_metrics.APP_UP.labels.return_value = mock_app_up
|
|
|
|
metadata: dict[str, str | None] = {'pod_id': 'test-pod-123', 'workflow_name': 'train_model'}
|
|
|
|
start_prometheus_server(mock_logger, metadata)
|
|
|
|
# Agora a mensagem é enviada via logger
|
|
mock_logger.custom_info.assert_called_once()
|
|
|
|
|
|
@patch('model_manager.worker.worker.start_http_server')
|
|
@patch('model_manager.worker.worker.metrics')
|
|
@patch('model_manager.worker.worker.os._exit')
|
|
def test_start_prometheus_server_prints_failure(
|
|
mock_exit, mock_metrics, mock_start_http_server, capsys, mock_env_vars, mock_logger
|
|
):
|
|
"""Test that start_prometheus_server prints failure message."""
|
|
from model_manager.worker.worker import start_prometheus_server
|
|
|
|
mock_start_http_server.side_effect = Exception('Test error')
|
|
|
|
metadata: dict[str, str | None] = {'pod_id': 'test-pod-123', 'workflow_name': 'train_model'}
|
|
|
|
start_prometheus_server(mock_logger, metadata)
|
|
|
|
# Agora o erro é logado via logger crítico
|
|
mock_logger.custom_critical.assert_called_once()
|
|
mock_exit.assert_called_once_with(1)
|