c5cd382350
- Updated `.env.example` to include new environment variables for MinIO and PyPI configuration. - Refactored `create_cleanup_schedule` to utilize runtime-specific task queues and improve schedule reconciliation logic. - Enhanced `Activities` class to require a default bucket in MinIO configuration. - Adjusted `requirements.txt` to specify version for `evidently`. - Updated tests to reflect changes in schedule creation and configuration handling.
896 lines
32 KiB
Python
896 lines
32 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
|
|
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')
|
|
async def test_main_successful_startup(
|
|
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')
|
|
async def test_main_handles_exception(
|
|
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')
|
|
async def test_main_temporal_client_configuration(
|
|
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')
|
|
async def test_main_worker_configuration(
|
|
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')
|
|
async def test_main_schedule_creation_failure_does_not_stop_worker(
|
|
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')
|
|
async def test_main_missing_runtime_uses_single_fallback(
|
|
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
|
|
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)
|