Code import - branch release/SIENTIAPDE-1646
This commit is contained in:
243
tests/laborious/worker/test_worker.py
Normal file
243
tests/laborious/worker/test_worker.py
Normal file
@@ -0,0 +1,243 @@
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from pytest import mark, raises
|
||||
|
||||
from laborious.worker import worker
|
||||
|
||||
|
||||
def _build_fake_activities():
|
||||
inst = MagicMock()
|
||||
inst.init_opc = MagicMock()
|
||||
inst.shutdown = MagicMock()
|
||||
inst.load_query_with_minio_offload = MagicMock()
|
||||
inst.retrain_model = MagicMock()
|
||||
inst.update_production_model = MagicMock()
|
||||
inst.format_retrain_report = MagicMock()
|
||||
inst.export_data_to_postgres = MagicMock()
|
||||
inst.load_custom_query = MagicMock()
|
||||
inst.calculate_simple_metrics = MagicMock()
|
||||
inst.get_reference_data = MagicMock()
|
||||
inst.calculate_drift = MagicMock()
|
||||
inst.request_predict = MagicMock()
|
||||
inst.request_transform = MagicMock()
|
||||
inst.input_gate = MagicMock()
|
||||
inst.mlflow_response_gate = MagicMock()
|
||||
inst.mlflow_content_gate = MagicMock()
|
||||
inst.format_transformed_data = MagicMock()
|
||||
inst.format_prediction = MagicMock()
|
||||
inst.format_default_prediction = MagicMock()
|
||||
inst.write_opc_data = MagicMock()
|
||||
inst.cleanup_minio_objects_expired = MagicMock()
|
||||
inst.repeat_last_prediction = MagicMock()
|
||||
inst.write_metrics = MagicMock()
|
||||
inst.write_pi_web_api_data = MagicMock()
|
||||
return inst
|
||||
|
||||
|
||||
def _build_fake_worker(async_result=None, async_error: Exception | None = None):
|
||||
w = MagicMock()
|
||||
|
||||
async def _run():
|
||||
if async_error is not None:
|
||||
raise async_error
|
||||
return async_result
|
||||
|
||||
w.run = MagicMock(side_effect=_run)
|
||||
return w
|
||||
|
||||
|
||||
@patch('laborious.worker.worker.start_http_server')
|
||||
def test_start_prometheus_server_success(mock_start_http):
|
||||
with patch.object(worker.metrics.APP_UP, 'labels') as labels:
|
||||
gauge = MagicMock()
|
||||
labels.return_value = gauge
|
||||
with patch('laborious.worker.worker.os.getenv', return_value='9090'):
|
||||
worker.start_prometheus_server()
|
||||
mock_start_http.assert_called_once_with(9090)
|
||||
gauge.set.assert_called_once_with(1)
|
||||
|
||||
|
||||
@patch('laborious.worker.worker.start_http_server', side_effect=RuntimeError('nope'))
|
||||
def test_start_prometheus_server_error_exits(_mock_start_http):
|
||||
with patch('laborious.worker.worker.os._exit', side_effect=SystemExit(1)) as m_exit:
|
||||
with raises(SystemExit):
|
||||
worker.start_prometheus_server()
|
||||
m_exit.assert_called_once_with(1)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_main_missing_runtime_exits_fast(monkeypatch):
|
||||
monkeypatch.setenv('RUNTIME', '')
|
||||
with (
|
||||
patch('laborious.worker.worker.start_prometheus_server'),
|
||||
patch('laborious.worker.worker.get_logger') as m_logger,
|
||||
patch('laborious.worker.worker.NotificationHandler'),
|
||||
patch('laborious.worker.worker.MetricsController'),
|
||||
patch.object(worker.metrics.APP_UP, 'labels') as labels,
|
||||
patch('laborious.worker.worker.sys.exit', side_effect=SystemExit(1)),
|
||||
):
|
||||
labels.return_value = MagicMock()
|
||||
with raises(SystemExit):
|
||||
await worker.main()
|
||||
assert m_logger.return_value.custom_critical.called
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_main_plugin_install_failure(monkeypatch):
|
||||
monkeypatch.setenv('RUNTIME', 'single')
|
||||
fake_activities = _build_fake_activities()
|
||||
fake_plugin = MagicMock()
|
||||
fake_plugin.install_runtime = AsyncMock(side_effect=RuntimeError('install failed'))
|
||||
with (
|
||||
patch('laborious.worker.worker.start_prometheus_server'),
|
||||
patch('laborious.worker.worker.get_logger'),
|
||||
patch(
|
||||
'laborious.worker.worker.build_mongodb_config',
|
||||
return_value={'connection_string': 'cs', 'database_name': 'db'},
|
||||
),
|
||||
patch('laborious.worker.worker.NotificationHandler'),
|
||||
patch('laborious.worker.worker.MetricsController'),
|
||||
patch(
|
||||
'laborious.worker.worker.build_plugin_store_config',
|
||||
return_value={
|
||||
'base_url': '',
|
||||
'owner': '',
|
||||
'repo': '',
|
||||
'username': None,
|
||||
'password': None,
|
||||
'branch': None,
|
||||
'cache_ttl_seconds': None,
|
||||
'pypi_index_url': '',
|
||||
'pypi_username': None,
|
||||
'pypi_password': None,
|
||||
},
|
||||
),
|
||||
patch('laborious.worker.worker.PluginStore', return_value=fake_plugin),
|
||||
patch('laborious.worker.worker.Activities', return_value=fake_activities),
|
||||
patch.object(worker.metrics.APP_UP, 'labels') as labels,
|
||||
patch('laborious.worker.worker.sys.exit', side_effect=SystemExit(1)),
|
||||
):
|
||||
labels.return_value = MagicMock()
|
||||
with raises(SystemExit):
|
||||
await worker.main()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_main_success_exit_zero(monkeypatch):
|
||||
monkeypatch.setenv('RUNTIME', 'single')
|
||||
fake_activities = _build_fake_activities()
|
||||
fake_plugin = MagicMock()
|
||||
fake_plugin.install_runtime = AsyncMock(return_value=None)
|
||||
fake_workers = [_build_fake_worker() for _ in range(4)]
|
||||
|
||||
with (
|
||||
patch('laborious.worker.worker.start_prometheus_server'),
|
||||
patch('laborious.worker.worker.get_logger'),
|
||||
patch(
|
||||
'laborious.worker.worker.build_mongodb_config',
|
||||
return_value={'connection_string': 'cs', 'database_name': 'db'},
|
||||
),
|
||||
patch('laborious.worker.worker.NotificationHandler') as m_notif_cls,
|
||||
patch('laborious.worker.worker.MetricsController'),
|
||||
patch(
|
||||
'laborious.worker.worker.build_plugin_store_config',
|
||||
return_value={
|
||||
'base_url': '',
|
||||
'owner': '',
|
||||
'repo': '',
|
||||
'username': None,
|
||||
'password': None,
|
||||
'branch': None,
|
||||
'cache_ttl_seconds': None,
|
||||
'pypi_index_url': '',
|
||||
'pypi_username': None,
|
||||
'pypi_password': None,
|
||||
},
|
||||
),
|
||||
patch('laborious.worker.worker.PluginStore', return_value=fake_plugin),
|
||||
patch('laborious.worker.worker.Activities', return_value=fake_activities),
|
||||
patch('laborious.worker.worker.build_postgres_config', return_value={}),
|
||||
patch('laborious.worker.worker.build_minio_config', return_value={}),
|
||||
patch('laborious.worker.worker.build_opc_config', return_value={}),
|
||||
patch('laborious.worker.worker.build_api_config', return_value={}),
|
||||
patch('laborious.worker.worker.PrometheusConfig', return_value=MagicMock()),
|
||||
patch('laborious.worker.worker.TelemetryConfig', return_value=MagicMock()),
|
||||
patch('laborious.worker.worker.Runtime', return_value=MagicMock()),
|
||||
patch(
|
||||
'laborious.worker.worker.client.Client.connect', new=AsyncMock(return_value=MagicMock())
|
||||
),
|
||||
patch('laborious.worker.worker.prepare_worker', side_effect=fake_workers) as m_prepare,
|
||||
patch.object(worker.metrics.APP_UP, 'labels') as labels,
|
||||
patch('laborious.worker.worker.sys.exit', side_effect=SystemExit(0)),
|
||||
):
|
||||
labels.return_value = MagicMock()
|
||||
with raises(SystemExit):
|
||||
await worker.main()
|
||||
|
||||
notif = m_notif_cls.return_value
|
||||
notif.shutdown.assert_called_once()
|
||||
fake_activities.shutdown.assert_called_once()
|
||||
assert m_prepare.call_count == 4
|
||||
prepare_calls = m_prepare.call_args_list
|
||||
assert prepare_calls[0].kwargs['runtime'] == 'single'
|
||||
assert prepare_calls[3].kwargs['runtime'] == 'single'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_main_worker_gather_error_exits_one(monkeypatch):
|
||||
monkeypatch.setenv('RUNTIME', 'single')
|
||||
fake_activities = _build_fake_activities()
|
||||
fake_plugin = MagicMock()
|
||||
fake_plugin.install_runtime = AsyncMock(return_value=None)
|
||||
fake_workers = [
|
||||
_build_fake_worker(async_error=RuntimeError('boom')),
|
||||
_build_fake_worker(),
|
||||
_build_fake_worker(),
|
||||
_build_fake_worker(),
|
||||
]
|
||||
|
||||
with (
|
||||
patch('laborious.worker.worker.start_prometheus_server'),
|
||||
patch('laborious.worker.worker.get_logger') as m_logger,
|
||||
patch(
|
||||
'laborious.worker.worker.build_mongodb_config',
|
||||
return_value={'connection_string': 'cs', 'database_name': 'db'},
|
||||
),
|
||||
patch('laborious.worker.worker.NotificationHandler'),
|
||||
patch('laborious.worker.worker.MetricsController'),
|
||||
patch(
|
||||
'laborious.worker.worker.build_plugin_store_config',
|
||||
return_value={
|
||||
'base_url': '',
|
||||
'owner': '',
|
||||
'repo': '',
|
||||
'username': None,
|
||||
'password': None,
|
||||
'branch': None,
|
||||
'cache_ttl_seconds': None,
|
||||
'pypi_index_url': '',
|
||||
'pypi_username': None,
|
||||
'pypi_password': None,
|
||||
},
|
||||
),
|
||||
patch('laborious.worker.worker.PluginStore', return_value=fake_plugin),
|
||||
patch('laborious.worker.worker.Activities', return_value=fake_activities),
|
||||
patch('laborious.worker.worker.build_postgres_config', return_value={}),
|
||||
patch('laborious.worker.worker.build_minio_config', return_value={}),
|
||||
patch('laborious.worker.worker.build_opc_config', return_value={}),
|
||||
patch('laborious.worker.worker.build_api_config', return_value={}),
|
||||
patch('laborious.worker.worker.PrometheusConfig', return_value=MagicMock()),
|
||||
patch('laborious.worker.worker.TelemetryConfig', return_value=MagicMock()),
|
||||
patch('laborious.worker.worker.Runtime', return_value=MagicMock()),
|
||||
patch(
|
||||
'laborious.worker.worker.client.Client.connect', new=AsyncMock(return_value=MagicMock())
|
||||
),
|
||||
patch('laborious.worker.worker.prepare_worker', side_effect=fake_workers),
|
||||
patch.object(worker.metrics.APP_UP, 'labels') as labels,
|
||||
patch('laborious.worker.worker.sys.exit', side_effect=SystemExit(1)),
|
||||
):
|
||||
labels.return_value = MagicMock()
|
||||
with raises(SystemExit):
|
||||
await worker.main()
|
||||
|
||||
assert m_logger.return_value.custom_error.called
|
||||
Reference in New Issue
Block a user