"""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'}