diff --git a/tests/worker/test_worker.py b/tests/worker/test_worker.py index 4980b0d..7a45d7f 100644 --- a/tests/worker/test_worker.py +++ b/tests/worker/test_worker.py @@ -528,6 +528,100 @@ async def test_main_worker_configuration( assert cleanup_call_args[1]['max_cached_workflows'] == 100 +@pytest.mark.asyncio +@patch('model_manager.worker.worker.create_cleanup_schedule') +@patch('model_manager.worker.worker.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.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_notification_handler_class, + mock_activities_class, + mock_runtime_class, + mock_client_class, + mock_worker_class, + 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 = AsyncMock() + mock_activities_class.return_value = mock_activities + + mock_runtime = Mock() + mock_runtime_class.return_value = mock_runtime + + mock_client_instance = AsyncMock() + mock_client_class.connect = AsyncMock(return_value=mock_client_instance) + + mock_worker_instance = Mock() + mock_worker_instance.run = AsyncMock(side_effect=asyncio.CancelledError()) + mock_worker_class.return_value = mock_worker_instance + + mock_app_up = Mock() + mock_metrics.APP_UP.labels.return_value = mock_app_up + + # Run main() - should not fail despite schedule creation error + with pytest.raises(SystemExit): + await main() + + # Verify schedule creation was attempted + mock_create_cleanup_schedule.assert_called_once() + + # Verify error was logged - check all custom_error calls + assert mock_logger.custom_error.call_count >= 1 + + # Find the call that contains the schedule error message + schedule_error_logged = False + for call in mock_logger.custom_error.call_args_list: + if 'Failed to configure cleanup schedule' in call[0][0]: + schedule_error_logged = True + break + + assert schedule_error_logged, 'Schedule creation error should be logged' + + # Verify workers were still created (startup continued) + assert mock_worker_class.call_count == 2 + + @patch('model_manager.worker.worker.asyncio.run') def test_main_entrypoint(mock_asyncio_run): """Test the __main__ entrypoint."""