|
|
|
|
@@ -30,6 +30,16 @@ def mock_env_vars():
|
|
|
|
|
yield env_vars
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
|
|
|
def runtime_ready(monkeypatch):
|
|
|
|
|
"""Boot as the re-executed process by default; the real os.execv must never run in tests."""
|
|
|
|
|
monkeypatch.setenv('SIENTIA_RUNTIME_READY', '1')
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class _StopBootError(Exception):
|
|
|
|
|
"""Stops main() at a chosen point, e.g. where os.execv would replace the process."""
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.fixture
|
|
|
|
|
def mock_logger():
|
|
|
|
|
"""Create a mock logger."""
|
|
|
|
|
@@ -844,6 +854,119 @@ async def test_main_missing_runtime_uses_single_fallback(
|
|
|
|
|
assert mock_prepare_worker.call_args_list[1].kwargs['runtime'] == 'single'
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
_PLUGIN_STORE_CONFIG = {
|
|
|
|
|
'base_url': 'http://sientia-plugin-store.svc.cluster.local',
|
|
|
|
|
'owner': 'sientia',
|
|
|
|
|
'repo': 'model-library-store',
|
|
|
|
|
'branch': 'main',
|
|
|
|
|
'username': 'gitea-user',
|
|
|
|
|
'password': 'gitea-password',
|
|
|
|
|
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
|
|
|
|
|
'pypi_username': None,
|
|
|
|
|
'pypi_password': None,
|
|
|
|
|
'cache_ttl_seconds': None,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
|
|
@patch('model_manager.worker.worker.RUNTIME', 'basic')
|
|
|
|
|
# Stops boot at Activities, so the red phase does not try to reach a real Temporal server.
|
|
|
|
|
@patch('model_manager.worker.worker.Activities', side_effect=_StopBootError)
|
|
|
|
|
@patch('model_manager.worker.worker.NotificationHandler')
|
|
|
|
|
@patch('model_manager.worker.worker.PluginStore')
|
|
|
|
|
@patch('model_manager.worker.worker.build_plugin_store_config', return_value=_PLUGIN_STORE_CONFIG)
|
|
|
|
|
@patch('model_manager.worker.worker.build_postgres_config', return_value={})
|
|
|
|
|
@patch('model_manager.worker.worker.build_mlflow_config', return_value={})
|
|
|
|
|
@patch('model_manager.worker.worker.build_minio_config', return_value={})
|
|
|
|
|
@patch(
|
|
|
|
|
'model_manager.worker.worker.build_mongodb_config',
|
|
|
|
|
return_value={'connection_string': 'mongodb://test', 'database_name': 'db', 'uri': 'u'},
|
|
|
|
|
)
|
|
|
|
|
@patch('model_manager.worker.worker.get_logger')
|
|
|
|
|
@patch('model_manager.worker.worker.start_prometheus_server')
|
|
|
|
|
@patch('model_manager.worker.worker.ensure_runtime_directories')
|
|
|
|
|
async def test_main_reexecs_after_installing_runtime(
|
|
|
|
|
_mock_dirs,
|
|
|
|
|
_mock_prometheus,
|
|
|
|
|
_mock_get_logger,
|
|
|
|
|
_mock_mongo,
|
|
|
|
|
_mock_minio,
|
|
|
|
|
_mock_mlflow,
|
|
|
|
|
_mock_postgres,
|
|
|
|
|
_mock_ps_config,
|
|
|
|
|
mock_plugin_store_class,
|
|
|
|
|
_mock_notification,
|
|
|
|
|
mock_activities_class,
|
|
|
|
|
monkeypatch,
|
|
|
|
|
):
|
|
|
|
|
"""First boot installs the runtime, then replaces the process so imports see it on disk."""
|
|
|
|
|
from model_manager.worker import worker as worker_module
|
|
|
|
|
|
|
|
|
|
monkeypatch.setenv('SIENTIA_RUNTIME_READY', '0')
|
|
|
|
|
mock_plugin_store_class.return_value.install_runtime = AsyncMock(return_value=None)
|
|
|
|
|
|
|
|
|
|
with (
|
|
|
|
|
patch('model_manager.worker.worker.os.execv', side_effect=_StopBootError) as m_execv,
|
|
|
|
|
patch.object(sys, 'orig_argv', ['python', '-m', 'model_manager.worker.worker']),
|
|
|
|
|
pytest.raises(_StopBootError),
|
|
|
|
|
):
|
|
|
|
|
await worker_module.main()
|
|
|
|
|
|
|
|
|
|
mock_plugin_store_class.return_value.install_runtime.assert_awaited_once_with(
|
|
|
|
|
runtime_name='basic'
|
|
|
|
|
)
|
|
|
|
|
assert os.environ['SIENTIA_RUNTIME_READY'] == '1'
|
|
|
|
|
m_execv.assert_called_once_with(
|
|
|
|
|
sys.executable, [sys.executable, '-m', 'model_manager.worker.worker']
|
|
|
|
|
)
|
|
|
|
|
mock_activities_class.assert_not_called()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
|
|
@patch('model_manager.worker.worker.RUNTIME', 'basic')
|
|
|
|
|
@patch('model_manager.worker.worker.Activities', side_effect=_StopBootError)
|
|
|
|
|
@patch('model_manager.worker.worker.NotificationHandler')
|
|
|
|
|
@patch('model_manager.worker.worker.PluginStore')
|
|
|
|
|
@patch('model_manager.worker.worker.build_plugin_store_config', return_value=_PLUGIN_STORE_CONFIG)
|
|
|
|
|
@patch('model_manager.worker.worker.build_postgres_config', return_value={})
|
|
|
|
|
@patch('model_manager.worker.worker.build_mlflow_config', return_value={})
|
|
|
|
|
@patch('model_manager.worker.worker.build_minio_config', return_value={})
|
|
|
|
|
@patch(
|
|
|
|
|
'model_manager.worker.worker.build_mongodb_config',
|
|
|
|
|
return_value={'connection_string': 'mongodb://test', 'database_name': 'db', 'uri': 'u'},
|
|
|
|
|
)
|
|
|
|
|
@patch('model_manager.worker.worker.get_logger')
|
|
|
|
|
@patch('model_manager.worker.worker.start_prometheus_server')
|
|
|
|
|
@patch('model_manager.worker.worker.ensure_runtime_directories')
|
|
|
|
|
async def test_main_skips_install_when_runtime_ready(
|
|
|
|
|
_mock_dirs,
|
|
|
|
|
_mock_prometheus,
|
|
|
|
|
_mock_get_logger,
|
|
|
|
|
_mock_mongo,
|
|
|
|
|
_mock_minio,
|
|
|
|
|
_mock_mlflow,
|
|
|
|
|
_mock_postgres,
|
|
|
|
|
_mock_ps_config,
|
|
|
|
|
mock_plugin_store_class,
|
|
|
|
|
_mock_notification,
|
|
|
|
|
mock_activities_class,
|
|
|
|
|
):
|
|
|
|
|
"""The re-executed process skips the install and the exec and goes on booting."""
|
|
|
|
|
from model_manager.worker import worker as worker_module
|
|
|
|
|
|
|
|
|
|
mock_plugin_store_class.return_value.install_runtime = AsyncMock(return_value=None)
|
|
|
|
|
|
|
|
|
|
with (
|
|
|
|
|
patch('model_manager.worker.worker.os.execv') as m_execv,
|
|
|
|
|
pytest.raises(_StopBootError),
|
|
|
|
|
):
|
|
|
|
|
await worker_module.main()
|
|
|
|
|
|
|
|
|
|
mock_plugin_store_class.return_value.install_runtime.assert_not_awaited()
|
|
|
|
|
m_execv.assert_not_called()
|
|
|
|
|
mock_activities_class.assert_called_once()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@patch('model_manager.worker.worker.asyncio.run')
|
|
|
|
|
def test_main_entrypoint(mock_asyncio_run):
|
|
|
|
|
"""Test the __main__ entrypoint."""
|
|
|
|
|
|