dfbdd4b605
The probes only check app_up, which was set to 1 when the metrics server started, before install_runtime. A pod that would fail the guard could turn Ready first, so a rollout would remove the old pod. app_up now goes to 1 right after the guard. Also: copy sys.modules keys before iterating, raise the sientia_model floor to the runtime's >=0.14.0, pin the same runtime versions in requirements-local.txt, test that the snapshot is taken before the install and the real importlib.metadata path, and document the boot failure in README and CLAUDE.md.
1008 lines
37 KiB
Python
1008 lines
37 KiB
Python
"""Unit tests for worker module."""
|
|
|
|
import asyncio
|
|
import os
|
|
import sys
|
|
from unittest.mock import AsyncMock, MagicMock, 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
|
|
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)
|
|
|
|
# app_up goes to 1 in main(), after the runtime version guard.
|
|
mock_app_up.set.assert_not_called()
|
|
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()
|
|
# Ready after the guard, down again on shutdown.
|
|
assert [c.args[0] for c in mock_app_up.set.call_args_list] == [1, 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'
|
|
|
|
|
|
@patch('model_manager.worker.worker.asyncio.run')
|
|
def test_main_entrypoint(mock_asyncio_run):
|
|
"""Test the __main__ entrypoint."""
|
|
# Import and execute the main block
|
|
import model_manager.worker.worker as worker_module
|
|
|
|
# patch.object restores main afterwards; assigning it leaked the mock into later tests.
|
|
with patch.object(sys, 'argv', ['worker.py']), patch.object(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)
|
|
|
|
|
|
def test_loaded_distribution_versions_reads_imported_packages():
|
|
"""An imported module maps to its distribution's installed version."""
|
|
import importlib.metadata
|
|
|
|
from model_manager.worker.worker import _loaded_distribution_versions
|
|
|
|
assert _loaded_distribution_versions()['pytest'] == importlib.metadata.version('pytest')
|
|
|
|
|
|
def test_runtime_changed_loaded_packages_reports_only_changes():
|
|
"""Only distributions whose version on disk differs after the install are reported."""
|
|
import importlib.metadata
|
|
|
|
from model_manager.worker.worker import _runtime_changed_loaded_packages
|
|
|
|
on_disk = {'scikit-learn': '1.8.0', 'numpy': '1.26.4'}
|
|
|
|
def fake_version(dist):
|
|
if dist not in on_disk:
|
|
raise importlib.metadata.PackageNotFoundError(dist)
|
|
return on_disk[dist]
|
|
|
|
before = {'scikit-learn': '1.9.1', 'numpy': '1.26.4', 'gone': '1.0'}
|
|
with patch('model_manager.worker.worker.importlib.metadata.version', side_effect=fake_version):
|
|
changed = _runtime_changed_loaded_packages(before)
|
|
|
|
assert changed == {'scikit-learn': '1.9.1 -> 1.8.0', 'gone': '1.0 -> removed'}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_main_runtime_changed_loaded_package_exits(mock_env_vars, mock_logger):
|
|
"""Boot stops before Activities when install_runtime changed an imported package."""
|
|
import model_manager.worker.worker as worker_module
|
|
|
|
plugin_store = Mock()
|
|
plugin_store.install_runtime = AsyncMock(return_value=None)
|
|
before = {'scikit-learn': '1.9.1'}
|
|
|
|
def snapshot():
|
|
# Taken after install_runtime, it would already see 1.8.0 and the guard would pass.
|
|
assert plugin_store.install_runtime.await_count == 0
|
|
return before
|
|
|
|
mongo = {'connection_string': 'mongodb://test', 'database_name': 'test_db', 'uri': 'x'}
|
|
with (
|
|
patch.object(worker_module, 'RUNTIME', 'basic'),
|
|
patch.object(worker_module, 'ensure_runtime_directories'),
|
|
patch.object(worker_module, 'metrics') as mock_metrics,
|
|
patch.object(worker_module, 'start_prometheus_server'),
|
|
patch.object(worker_module, 'get_logger', return_value=mock_logger),
|
|
patch.object(worker_module, 'build_mongodb_config', return_value=mongo),
|
|
patch.object(worker_module, 'NotificationHandler'),
|
|
patch.object(worker_module, 'build_plugin_store_config', return_value=MagicMock()),
|
|
patch.object(worker_module, 'PluginStore', return_value=plugin_store),
|
|
patch.object(worker_module, '_loaded_distribution_versions', side_effect=snapshot),
|
|
patch.object(
|
|
worker_module,
|
|
'_runtime_changed_loaded_packages',
|
|
return_value={'scikit-learn': '1.9.1 -> 1.8.0'},
|
|
) as mock_changed,
|
|
patch.object(worker_module, 'Activities') as mock_activities_class,
|
|
patch.object(worker_module.sys, 'exit', side_effect=SystemExit(1)) as mock_exit,
|
|
):
|
|
with pytest.raises(SystemExit):
|
|
await worker_module.main()
|
|
|
|
mock_exit.assert_called_once_with(1)
|
|
mock_activities_class.assert_not_called()
|
|
plugin_store.install_runtime.assert_awaited_once()
|
|
mock_changed.assert_called_once_with(before)
|
|
mock_metrics.APP_UP.labels.return_value.set.assert_called_once_with(0)
|
|
message = mock_logger.custom_critical.call_args.args[0]
|
|
assert "'scikit-learn': '1.9.1 -> 1.8.0'" in message
|
|
|
|
|
|
def test_guard_detects_version_change_on_disk(tmp_path, monkeypatch):
|
|
"""The real importlib.metadata path: a package swapped on disk after import is reported."""
|
|
import types
|
|
|
|
from model_manager.worker.worker import (
|
|
_loaded_distribution_versions,
|
|
_runtime_changed_loaded_packages,
|
|
)
|
|
|
|
info = tmp_path / 'fakepkg_2214-1.0.dist-info'
|
|
info.mkdir()
|
|
(info / 'METADATA').write_text('Metadata-Version: 2.1\nName: fakepkg-2214\nVersion: 1.0\n')
|
|
(info / 'top_level.txt').write_text('fakepkg_2214\n')
|
|
monkeypatch.syspath_prepend(str(tmp_path))
|
|
# The guard only reads sys.modules keys; setitem removes the entry again on teardown.
|
|
monkeypatch.setitem(sys.modules, 'fakepkg_2214', types.ModuleType('fakepkg_2214'))
|
|
|
|
before = _loaded_distribution_versions()
|
|
assert before['fakepkg-2214'] == '1.0'
|
|
|
|
# What pip does on an upgrade: the old dist-info goes away and a new one takes its place.
|
|
new_info = info.rename(tmp_path / 'fakepkg_2214-2.0.dist-info')
|
|
(new_info / 'METADATA').write_text('Metadata-Version: 2.1\nName: fakepkg-2214\nVersion: 2.0\n')
|
|
|
|
assert _runtime_changed_loaded_packages(before) == {'fakepkg-2214': '1.0 -> 2.0'}
|