QTZPOC-20
QTZPOC-20: Enhance system robustness and test coverage across bundle processing, activity error handling, and worker operations. * Strengthen bundle reader validation with new tests for archive integrity, security vulnerabilities (e.g., path traversal, duplicate entries, special file types), and I/O error resilience. * Improve error handling and status propagation in PI Web API, MLflow, and OPC activities, including partial write scenarios and OPC session management. * Add comprehensive tests for worker startup, configuration, graceful shutdown, and Prometheus metrics server initialization. * Introduce a new utility for intelligent DataFrame debugging, summarizing large dataframes for concise logging, and refine exception chain summarization. * Update code coverage configuration to exclude @overload decorator lines.
This commit is contained in:
@@ -179,6 +179,7 @@ exclude_lines = [
|
||||
"if TYPE_CHECKING:",
|
||||
"class .*\\bProtocol\\):",
|
||||
"@(abc\\.)?abstractmethod",
|
||||
"@overload",
|
||||
]
|
||||
|
||||
[tool.coverage.html]
|
||||
|
||||
@@ -23,6 +23,7 @@ import hashlib
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import stat
|
||||
import struct
|
||||
import unicodedata
|
||||
import zipfile
|
||||
@@ -118,6 +119,15 @@ class ZipSpec:
|
||||
traversal_entry: str | None = None
|
||||
bomb_bytes: int = 0
|
||||
extra_files: dict[str, bytes] = field(default_factory=dict)
|
||||
# name -> a raw `stat.S_IF*` mode, for an entry that is neither a symlink, a regular file nor a
|
||||
# directory (e.g. `stat.S_IFIFO`) — gate 4 must reject these on their file-type bits alone.
|
||||
special_mode_entries: dict[str, int] = field(default_factory=dict)
|
||||
# Directory entries the producer never writes (`build_zip` stores none), but a gate 5 defense
|
||||
# this repository still carries has to be exercised by something.
|
||||
directory_entries: list[str] = field(default_factory=list)
|
||||
# A file name already in `files`/`extra_files`, written a second time under the same name — two
|
||||
# archive entries claiming the same path, which only gate 5's extraction can catch.
|
||||
repeat_file: str | None = None
|
||||
|
||||
|
||||
def build_zip(spec: ZipSpec | None = None) -> bytes:
|
||||
@@ -145,6 +155,19 @@ def build_zip(spec: ZipSpec | None = None) -> bytes:
|
||||
archive.writestr(_regular_file_info(spec.traversal_entry), b'traversal\n')
|
||||
if spec.bomb_bytes:
|
||||
archive.writestr(_regular_file_info('artifacts/bomb.csv'), b'\0' * spec.bomb_bytes)
|
||||
for name, mode in spec.special_mode_entries.items():
|
||||
info = zipfile.ZipInfo(name)
|
||||
info.create_system = 3
|
||||
info.external_attr = mode << 16
|
||||
archive.writestr(info, b'')
|
||||
for name in spec.directory_entries:
|
||||
info = zipfile.ZipInfo(name if name.endswith('/') else f'{name}/')
|
||||
info.create_system = 3
|
||||
info.external_attr = (stat.S_IFDIR | 0o755) << 16
|
||||
archive.writestr(info, b'')
|
||||
if spec.repeat_file:
|
||||
content = {**spec.files, **spec.extra_files}[spec.repeat_file]
|
||||
archive.writestr(_regular_file_info(spec.repeat_file), content)
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
|
||||
@@ -220,6 +220,44 @@ async def test_write_pi_web_api_data_success(mock_dataframe, api, base_input_dat
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.api.DataFrame')
|
||||
async def test_write_pi_web_api_data_partial_prediction_mismatch_downgrades_confidence(
|
||||
mock_dataframe, api, base_input_data
|
||||
):
|
||||
"""A partial write (some tags report Errors) doesn't raise — it's caught by the tag-count
|
||||
mismatch check, which still must downgrade confidence and attach the mismatch message."""
|
||||
input_data = {
|
||||
**base_input_data,
|
||||
'pi_web_api_output_config': {
|
||||
'endpoint': 'https://test-pi-server.com/piwebapi',
|
||||
'prediction_tags': {'tag1': 'web_id_1', 'tag2': 'web_id_2'},
|
||||
'confidence_tags': {'tag3': 'web_id_3'},
|
||||
},
|
||||
}
|
||||
mock_df = _create_mock_dataframe()
|
||||
mock_dataframe.return_value = mock_df
|
||||
|
||||
api.pi_web_api_client.write_value.side_effect = [
|
||||
[{'WebId': 'web_id_1', 'Errors': []}, {'WebId': 'web_id_2', 'Errors': ['bad tag']}],
|
||||
[{'WebId': 'web_id_3', 'Errors': []}],
|
||||
]
|
||||
|
||||
await api.write_pi_web_api_data(input_data)
|
||||
|
||||
mock_df.__setitem__.assert_any_call(
|
||||
'prediction_confidence', PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
||||
)
|
||||
mock_df.__setitem__.assert_any_call('comments', ANY)
|
||||
api.send_notification_async.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
|
||||
message=ANY,
|
||||
block='write_pi_web_api_data',
|
||||
level=NotificationLevel.ERROR,
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.api.DataFrame')
|
||||
async def test_write_pi_web_api_data_prediction_error(mock_dataframe, api, base_input_data):
|
||||
|
||||
@@ -222,6 +222,39 @@ async def test_request_predict(mock_to_datetime, mock_from_dataframe, mlflow):
|
||||
assert response_data == mock_from_dataframe.return_value
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch(
|
||||
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
@patch('laborious.activities.mlflow.to_datetime')
|
||||
async def test_request_predict_skips_timestamp_reinsertion_when_configured(
|
||||
mock_to_datetime, mock_from_dataframe, mlflow
|
||||
):
|
||||
"""skip_predict_timestamp=True leaves the caller's index alone — no reinsertion round-trip."""
|
||||
data_mock = MagicMock()
|
||||
payload = AsyncMock()
|
||||
payload.retrieve = AsyncMock(return_value=data_mock)
|
||||
|
||||
input_data = {
|
||||
**metadata,
|
||||
'data': payload,
|
||||
'model_name': 'test_model',
|
||||
'model_config': {'skip_predict_timestamp': True},
|
||||
}
|
||||
|
||||
predict_response = {'success': True, 'content': MagicMock()}
|
||||
mlflow.model_monitoring_repository.predict.return_value = predict_response
|
||||
|
||||
await mlflow.request_predict(input_data)
|
||||
|
||||
data_mock.replace.assert_called_once_with(np.nan, None, inplace=True)
|
||||
mock_to_datetime.assert_not_called()
|
||||
mlflow.model_monitoring_repository.predict.assert_called_once_with(
|
||||
'test_model', data_mock, {'skip_predict_timestamp': True}, metadata['metadata']
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch(
|
||||
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
|
||||
|
||||
@@ -166,6 +166,17 @@ def upload_input(extracted_dir: Path, **overrides) -> dict:
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_upload_artifacts_requires_the_mlflow_repository(
|
||||
import_activity: ModelImport, tmp_path: Path
|
||||
):
|
||||
import_activity.import_mlflow_repository = None
|
||||
payload = upload_input(tmp_path / 'extract')
|
||||
|
||||
with raises(ValueError, match='MLflow repository not initialized'):
|
||||
await import_activity.upload_import_artifacts(payload)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_every_import_creates_a_new_run_with_the_recorded_name(
|
||||
import_activity: ModelImport, mlflow_client: MagicMock, tmp_path: Path
|
||||
@@ -235,6 +246,15 @@ async def test_the_run_is_closed_when_its_tree_is_in(
|
||||
# ------------------------------------------------------------------------ registration
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_register_version_requires_the_mlflow_repository(import_activity: ModelImport):
|
||||
import_activity.import_mlflow_repository = None
|
||||
payload = base(model_name='factory_model', run_id=RUN_ID)
|
||||
|
||||
with raises(ValueError, match='MLflow repository not initialized'):
|
||||
await import_activity.register_import_model_version(payload)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_version_source_is_the_shape_the_legacy_loader_parses(
|
||||
import_activity: ModelImport, mlflow_client: MagicMock
|
||||
@@ -511,6 +531,17 @@ async def test_a_missing_mongo_repository_is_a_configuration_error(
|
||||
await import_activity.write_import_model_document(payload)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_next_model_document_id_also_guards_a_missing_mongo_repository(
|
||||
import_activity: ModelImport,
|
||||
):
|
||||
"""The helper repeats the guard so it stays safe if ever called from somewhere new."""
|
||||
import_activity.mongo_repository = None
|
||||
|
||||
with raises(ValueError, match='MongoDB repository not initialized'):
|
||||
await import_activity._next_model_document_id('models', {})
|
||||
|
||||
|
||||
def test_no_import_module_writes_an_active_document():
|
||||
"""Asserted from the source too, so `active: false` cannot be flipped by a later edit."""
|
||||
source = Path('laborious/activities/model_import.py').read_text()
|
||||
@@ -577,6 +608,15 @@ async def test_cleanup_never_touches_the_uploaded_object(
|
||||
client.remove_objects.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_ensure_retention_requires_the_minio_repository(import_activity: ModelImport):
|
||||
import_activity.import_minio_repository = None
|
||||
payload = base()
|
||||
|
||||
with raises(ValueError, match='Minio repository not initialized'):
|
||||
await import_activity.ensure_import_bundle_retention(payload)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_expiry_rule_is_applied_to_the_import_prefix(import_activity: ModelImport):
|
||||
client = import_activity.import_minio_repository.minio_client
|
||||
|
||||
@@ -221,6 +221,18 @@ def test_the_detail_records_a_cleanup_failure_without_a_code():
|
||||
assert 'code' not in detail
|
||||
|
||||
|
||||
def test_the_detail_carries_the_bundle_metadata_as_source():
|
||||
detail = json.loads(
|
||||
build_import_detail(workflow_id='wf-1', source={'exported_by': 'user@example.com'})
|
||||
)['import']
|
||||
assert detail['source'] == {'exported_by': 'user@example.com'}
|
||||
|
||||
|
||||
def test_the_detail_omits_source_when_empty():
|
||||
detail = json.loads(build_import_detail(workflow_id='wf-1', source={}))['import']
|
||||
assert 'source' not in detail
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------ the claim
|
||||
|
||||
|
||||
@@ -834,3 +846,13 @@ async def test_the_engine_is_disposed_on_shutdown(import_activity: ModelImport,
|
||||
import_activity.close()
|
||||
|
||||
assert engine.disposed is True
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_close_is_a_no_op_without_an_engine(import_activity: ModelImport, engine):
|
||||
"""A partially-constructed instance (or a second close) must not blow up on a missing engine."""
|
||||
del import_activity.status_engine
|
||||
|
||||
import_activity.close()
|
||||
|
||||
assert engine.disposed is False
|
||||
|
||||
@@ -594,6 +594,39 @@ async def test_write_opc_data_success(mock_dataframe, opc):
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.opc.DataFrame')
|
||||
async def test_write_opc_data_propagates_session_bad_and_reconnect_in_progress(mock_dataframe, opc):
|
||||
"""A single server reporting session_bad/reconnect_in_progress must flip the aggregated
|
||||
flags passed to process_confidence, even though nothing raised."""
|
||||
input_data = {
|
||||
**metadata,
|
||||
'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
|
||||
'opc_output_config': {
|
||||
'server1': {
|
||||
'prediction_tags': {'tag1': {'data_type': 'float'}},
|
||||
'confidence_tags': {'tag2': {'data_type': 'float'}},
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
opc.manage_output_tags = AsyncMock(
|
||||
return_value=(False, {'tag1': 0.1, 'tag2': 0.2}, True, 'BadSessionClosed', True)
|
||||
)
|
||||
opc.process_confidence = MagicMock(return_value={'data': 'data'})
|
||||
|
||||
await opc.write_opc_data(input_data)
|
||||
|
||||
opc.process_confidence.assert_called_once_with(
|
||||
mock_dataframe.return_value,
|
||||
False,
|
||||
metadata['metadata'],
|
||||
session_bad=True,
|
||||
opc_status='BadSessionClosed',
|
||||
reconnect_in_progress=True,
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_write_opc_data_empty_config(opc):
|
||||
# Arrange
|
||||
|
||||
@@ -1,7 +1,11 @@
|
||||
"""The seven gates: what each one rejects, in what order, and what it leaves on disk."""
|
||||
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import stat
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
|
||||
import nacl.bindings as sodium
|
||||
@@ -173,6 +177,16 @@ def test_gate_2_rejects_a_truncated_transfer(work_dir, reader):
|
||||
assert error.value.gate == 2
|
||||
|
||||
|
||||
def test_gate_2_wraps_an_os_error_reading_the_file(work_dir, reader):
|
||||
"""A download that never landed on disk is an OS error, not a format one."""
|
||||
missing = work_dir / 'never-downloaded.sientia'
|
||||
|
||||
with pytest.raises(BundleGateError, match='could not be read') as error:
|
||||
reader.open(missing, factory.DEFAULT_PASSWORD, work_dir / 'extract', work_dir)
|
||||
|
||||
assert error.value.gate == 2
|
||||
|
||||
|
||||
def _never_called(message: str):
|
||||
def _raise(*_args, **_kwargs):
|
||||
raise AssertionError(message)
|
||||
@@ -220,6 +234,22 @@ def test_gate_3_stream_without_a_final_tag_is_rejected(work_dir, reader):
|
||||
assert AMBIGUOUS_REASON in str(error.value)
|
||||
|
||||
|
||||
def test_gate_3_wraps_an_os_error_while_writing_the_decrypted_zip(work_dir, reader, monkeypatch):
|
||||
"""A write failure while streaming the plaintext to disk is reported as disk trouble, not as
|
||||
a bad password — the AEAD never even got to disagree."""
|
||||
payload = factory.make_bundle()
|
||||
|
||||
def _raise_os_error(*_args, **_kwargs):
|
||||
raise OSError('disk full')
|
||||
|
||||
monkeypatch.setattr(reader, '_pull_stream', _raise_os_error)
|
||||
|
||||
with pytest.raises(BundleGateError, match='could not be written') as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 3
|
||||
|
||||
|
||||
def test_gate_3_leaves_no_decrypted_zip_behind(work_dir, reader):
|
||||
payload = factory.make_bundle()
|
||||
|
||||
@@ -316,6 +346,27 @@ def test_gate_4_rejects_content_that_is_not_an_archive(work_dir, reader):
|
||||
assert error.value.gate == 4
|
||||
|
||||
|
||||
def test_gate_4_rejects_an_empty_archive(work_dir, reader):
|
||||
payload = factory.make_bundle(factory.ZipSpec(files={}, include_metadata=False))
|
||||
|
||||
with pytest.raises(BundleGateError, match='archive is empty') as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 4
|
||||
|
||||
|
||||
def test_gate_4_rejects_an_entry_that_is_neither_a_file_nor_a_directory(work_dir, reader):
|
||||
"""A FIFO's mode bits are neither a symlink's, a regular file's, nor a directory's."""
|
||||
payload = factory.make_bundle(
|
||||
factory.ZipSpec(special_mode_entries={'artifacts/pipe': stat.S_IFIFO | 0o644})
|
||||
)
|
||||
|
||||
with pytest.raises(BundleGateError, match='not a regular file') as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 4
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- gate 5
|
||||
|
||||
|
||||
@@ -338,6 +389,61 @@ def test_gate_5_extraction_writes_plain_files(work_dir, reader):
|
||||
assert path.stat().st_mode & 0o111 == 0, path
|
||||
|
||||
|
||||
def test_gate_5_extraction_creates_declared_directory_entries(work_dir, reader):
|
||||
"""The producer never writes directory entries, but a bundle from anywhere else might."""
|
||||
payload = factory.make_bundle(factory.ZipSpec(directory_entries=['artifacts/empty_dir']))
|
||||
|
||||
opened = _open(reader, work_dir, payload)
|
||||
|
||||
assert (opened.extracted_dir / 'artifacts' / 'empty_dir').is_dir()
|
||||
|
||||
|
||||
def test_gate_5_rejects_a_duplicate_entry_name(work_dir, reader):
|
||||
"""Two archive entries claiming the same path: gate 4 checks each name on its own and cannot
|
||||
see the collision, so extraction is where an overwrite is caught — and rejected."""
|
||||
payload = factory.make_bundle(factory.ZipSpec(repeat_file='artifacts/prediction_model/MLmodel'))
|
||||
|
||||
with pytest.raises(BundleGateError, match='overwrite') as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 5
|
||||
assert not (work_dir / 'extract').exists()
|
||||
|
||||
|
||||
def test_gate_5_wraps_an_os_error_during_unpacking(work_dir, reader, monkeypatch):
|
||||
"""A write failure partway through extraction (e.g. disk full) is reported as unpacking
|
||||
trouble, not surfaced as the raw OSError."""
|
||||
payload = factory.make_bundle()
|
||||
|
||||
def _raise_os_error(*_args, **_kwargs):
|
||||
raise OSError('disk full')
|
||||
|
||||
monkeypatch.setattr('laborious.utils.bundle.reader.shutil.copyfileobj', _raise_os_error)
|
||||
|
||||
with pytest.raises(BundleGateError, match='could not be unpacked') as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 5
|
||||
assert not (work_dir / 'extract').exists()
|
||||
|
||||
|
||||
def test_extract_member_rechecks_containment_on_the_resolved_path(work_dir, reader):
|
||||
"""Gate 4 already rejects a traversal by name; extraction re-checks on the *resolved* path as
|
||||
a second, independent line of defense."""
|
||||
root = work_dir / 'root'
|
||||
root.mkdir()
|
||||
buffer = io.BytesIO()
|
||||
with zipfile.ZipFile(buffer, 'w') as archive:
|
||||
archive.writestr('safe.txt', b'hello')
|
||||
|
||||
with zipfile.ZipFile(io.BytesIO(buffer.getvalue())) as archive:
|
||||
entry = archive.infolist()[0]
|
||||
with pytest.raises(BundleGateError, match='resolves outside') as error:
|
||||
reader._extract_member(archive, entry, root, work_dir / 'somewhere-else')
|
||||
|
||||
assert error.value.gate == 5
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- gate 6
|
||||
|
||||
|
||||
@@ -375,6 +481,10 @@ def test_gate_5_extraction_writes_plain_files(work_dir, reader):
|
||||
),
|
||||
'experiment name',
|
||||
),
|
||||
(
|
||||
factory.ZipSpec(document={'metadata': 'not-a-dict', 'parameters': {}}),
|
||||
'not objects',
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_gate_6_rejects_a_malformed_metadata_document(work_dir, reader, spec, match):
|
||||
@@ -457,6 +567,19 @@ def test_gate_7_rejects_content_a_model_export_should_not_carry(work_dir, reader
|
||||
assert not (work_dir / 'extract').exists()
|
||||
|
||||
|
||||
def test_gate_7_rejects_a_non_regular_file_left_in_the_tree(work_dir, reader):
|
||||
"""Extraction itself only ever writes plain files and directories; this is the check that
|
||||
would catch anything else showing up in the tree by some other means."""
|
||||
root = work_dir / 'weird_root'
|
||||
root.mkdir()
|
||||
os.mkfifo(root / 'pipe')
|
||||
|
||||
with pytest.raises(BundleGateError, match='not a file') as error:
|
||||
reader._gate_7_content_policy(root)
|
||||
|
||||
assert error.value.gate == 7
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ ordering
|
||||
|
||||
|
||||
|
||||
@@ -536,6 +536,40 @@ def test_a_failure_that_arrives_as_nothing_says_so_rather_than_guessing():
|
||||
assert summary == ''
|
||||
|
||||
|
||||
def test_summarize_failure_bounds_a_cyclic_cause_chain():
|
||||
"""A cycle never hits a `None` cause, so the walk must stop by exhausting its own bound."""
|
||||
first = RuntimeError('first')
|
||||
second = RuntimeError('second')
|
||||
first.__cause__ = second
|
||||
second.__cause__ = first
|
||||
|
||||
kind, _summary = summarize_failure(first)
|
||||
|
||||
assert kind == 'RuntimeError'
|
||||
|
||||
|
||||
def test_a_leading_blank_line_is_skipped_for_the_first_real_one():
|
||||
class BlankFirstLineError(Exception):
|
||||
def __str__(self) -> str:
|
||||
return '\n \nreal message'
|
||||
|
||||
_kind, summary = summarize_failure(BlankFirstLineError())
|
||||
|
||||
assert summary == 'real message'
|
||||
|
||||
|
||||
def test_describe_failure_chain_reports_a_message_less_link_by_its_type_alone():
|
||||
"""`Exception()` has no message at all — the link still has to be readable on its own."""
|
||||
inner = Exception()
|
||||
outer = RuntimeError('outer failure')
|
||||
outer.__cause__ = inner
|
||||
|
||||
kind, detail = describe_failure_chain(outer)
|
||||
|
||||
assert kind == 'Exception'
|
||||
assert detail.splitlines() == ['Exception', 'RuntimeError: outer failure']
|
||||
|
||||
|
||||
# ------------------------------------------------- the channel is the only reporter
|
||||
|
||||
|
||||
|
||||
@@ -214,6 +214,27 @@ async def test_from_dataframe_offloaded(mock_now):
|
||||
minio.upload_file.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 10**9)
|
||||
async def test_from_dataframe_inline_keeps_a_provided_last_timestamp():
|
||||
"""An explicit last_timestamp must survive as-is, skipping the max() derivation."""
|
||||
minio = AsyncMock()
|
||||
logger = MagicMock()
|
||||
df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
|
||||
|
||||
result = await MinioDataFramePayload.from_dataframe(
|
||||
dataframe=df,
|
||||
minio_repo=minio,
|
||||
model_name='m',
|
||||
operation='initial',
|
||||
last_timestamp='2030-01-01T00:00:00+00:00',
|
||||
logger=logger,
|
||||
)
|
||||
|
||||
assert result.last_timestamp == '2030-01-01T00:00:00+00:00'
|
||||
logger.custom_debug.assert_called()
|
||||
|
||||
|
||||
def test_from_dict_inline():
|
||||
raw = {
|
||||
'last_timestamp': '2024-01-01T00:00:00+00:00',
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from datetime import datetime
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, Mock, patch
|
||||
|
||||
@@ -8,11 +9,13 @@ from asyncua.crypto.security_policies import SecurityPolicyBasic256
|
||||
from asyncua.ua.uaerrors import BadNodeIdUnknown, BadSessionIdInvalid
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
|
||||
from laborious import metrics
|
||||
from laborious.utils.repository.opc_repository import (
|
||||
OpcClientAlreadyExistsError,
|
||||
OpcClientNotInitializedError,
|
||||
OpcRepository,
|
||||
OpcSessionAlreadyConnectedError,
|
||||
_opc_authentication_token_str,
|
||||
is_reconnectable_opcua_bad,
|
||||
)
|
||||
|
||||
@@ -80,6 +83,40 @@ def test_init(opc_repository):
|
||||
assert opc_repository.last_reconnection_time is None
|
||||
|
||||
|
||||
def test_opc_authentication_token_str_no_protocol():
|
||||
client = MagicMock()
|
||||
client.uaclient.protocol = None
|
||||
|
||||
assert _opc_authentication_token_str(client) == 'unknown'
|
||||
|
||||
|
||||
def test_opc_authentication_token_str_no_token():
|
||||
client = MagicMock()
|
||||
client.uaclient.protocol = MagicMock(authentication_token=None)
|
||||
|
||||
assert _opc_authentication_token_str(client) == 'unknown'
|
||||
|
||||
|
||||
def test_opc_authentication_token_str_swallows_errors():
|
||||
class _ExplodingClient:
|
||||
@property
|
||||
def uaclient(self):
|
||||
raise RuntimeError('boom')
|
||||
|
||||
assert _opc_authentication_token_str(_ExplodingClient()) == 'unknown'
|
||||
|
||||
|
||||
def test_is_session_open_swallows_errors(opc_repository):
|
||||
class _ExplodingClient:
|
||||
@property
|
||||
def uaclient(self):
|
||||
raise RuntimeError('boom')
|
||||
|
||||
opc_repository.client = _ExplodingClient()
|
||||
|
||||
assert opc_repository._is_session_open() is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_security(opc_repository, mock_client):
|
||||
opc_repository.client = mock_client
|
||||
@@ -160,6 +197,47 @@ async def test_create_client_raises_when_client_exists(opc_repository, mock_clie
|
||||
await opc_repository._create_client()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_client_builds_a_secured_client(opc_repository, mock_client):
|
||||
opc_repository.pod_id = 'pod-1'
|
||||
opc_repository.set_security = AsyncMock()
|
||||
|
||||
await opc_repository._create_client()
|
||||
|
||||
assert opc_repository.client is mock_client
|
||||
assert mock_client.name == 'pod-1'
|
||||
assert mock_client.application_name == 'pod-1'
|
||||
assert mock_client.application_uri == 'pod:1'
|
||||
assert mock_client.product_uri == 'pod:1'
|
||||
opc_repository.set_security.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_client_skips_security_without_a_cert_path(opc_repository, mock_client):
|
||||
opc_repository.cert_path = None
|
||||
opc_repository.set_security = AsyncMock()
|
||||
|
||||
await opc_repository._create_client()
|
||||
|
||||
assert opc_repository.client is mock_client
|
||||
opc_repository.set_security.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connect_locked_reuses_an_existing_client(opc_repository, mock_client):
|
||||
"""A client already set (e.g. surviving a prior failed open) is not recreated."""
|
||||
opc_repository.client = mock_client
|
||||
opc_repository._is_session_open = MagicMock(return_value=False)
|
||||
opc_repository._create_client = AsyncMock()
|
||||
opc_repository._open_session = AsyncMock(return_value=(True, {}))
|
||||
|
||||
result = await opc_repository._connect_locked()
|
||||
|
||||
opc_repository._create_client.assert_not_called()
|
||||
opc_repository._open_session.assert_called_once()
|
||||
assert result == (True, {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_open_session_success(opc_repository):
|
||||
closed_proto = MagicMock()
|
||||
@@ -435,6 +513,82 @@ async def test_write_data_write_value_failed(opc_repository, mock_client):
|
||||
assert error_data['attachment_content'] is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_data_with_empty_metadata_uses_unknown_labels(opc_repository, mock_client):
|
||||
"""No model_id/model_name in the write metadata: the metric labels fall back to 'unknown'."""
|
||||
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
|
||||
opc_repository.client = mock_client
|
||||
mock_node = AsyncMock()
|
||||
mock_client.get_node = MagicMock(return_value=mock_node)
|
||||
|
||||
result = await opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', {})
|
||||
|
||||
assert result == (True, {'response_time': ANY})
|
||||
last_call = opc_repository.emit_metric.call_args_list[-1]
|
||||
assert last_call.args[0] is metrics.OPC_WRITE_ATTEMPTS_TOTAL
|
||||
assert last_call.args[1]['model_id'] == 'unknown'
|
||||
assert last_call.args[1]['model_name'] == 'unknown'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_data_get_node_bad_session_id_schedules_reconnect(opc_repository, mock_client):
|
||||
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
|
||||
opc_repository.client = mock_client
|
||||
opc_repository._start_reconnect = AsyncMock()
|
||||
mock_client.get_node = MagicMock(side_effect=BadSessionIdInvalid())
|
||||
|
||||
is_success, error_data = await opc_repository.write_data(
|
||||
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
|
||||
)
|
||||
|
||||
opc_repository._start_reconnect.assert_called_once()
|
||||
assert is_success is False
|
||||
assert error_data['opc_error_kind'] == 'session_bad'
|
||||
assert error_data['opc_status'] == 'BadSessionIdInvalid'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_data_logs_inter_arrival_and_flags_a_stale_session(opc_repository, mock_client):
|
||||
"""The second write on the same session logs the gap since the previous one, and flags it
|
||||
when the gap outlives the session's revised timeout."""
|
||||
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
|
||||
opc_repository.client = mock_client
|
||||
mock_client.session_timeout = 0.001 # 1 microsecond: any real gap will exceed it
|
||||
mock_node = AsyncMock()
|
||||
mock_client.get_node = MagicMock(return_value=mock_node)
|
||||
|
||||
await opc_repository.write_data('ns=2;s=TestNode', 1.0, 'float', metadata['metadata'])
|
||||
await opc_repository.write_data('ns=2;s=TestNode', 2.0, 'float', metadata['metadata'])
|
||||
|
||||
logged = [c.args[0] for c in opc_repository.info.call_args_list]
|
||||
assert any('OPC write inter-arrival' in message for message in logged)
|
||||
emitted_metrics = [c.args[0] for c in opc_repository.emit_metric.call_args_list]
|
||||
assert metrics.OPC_WRITE_INTER_ARRIVAL_OVER_SESSION_TIMEOUT_TOTAL in emitted_metrics
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_write_inter_arrival_skips_the_metric_without_a_client(opc_repository):
|
||||
opc_repository._last_write_mono = time.monotonic() - 1
|
||||
opc_repository.client = None
|
||||
|
||||
await opc_repository._log_write_inter_arrival('session-1', 'ns=2;s=Node')
|
||||
|
||||
opc_repository.emit_metric.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_write_inter_arrival_skips_the_metric_within_the_session_timeout(
|
||||
opc_repository, mock_client
|
||||
):
|
||||
opc_repository._last_write_mono = time.monotonic()
|
||||
opc_repository.client = mock_client
|
||||
mock_client.session_timeout = 600_000 # 600s: comfortably longer than this test takes
|
||||
|
||||
await opc_repository._log_write_inter_arrival('session-1', 'ns=2;s=Node')
|
||||
|
||||
opc_repository.emit_metric.assert_not_called()
|
||||
|
||||
|
||||
def test_is_reconnectable_opcua_bad():
|
||||
assert is_reconnectable_opcua_bad(BadSessionIdInvalid()) is True
|
||||
assert is_reconnectable_opcua_bad(BadNodeIdUnknown()) is False
|
||||
@@ -491,6 +645,22 @@ async def test_start_reconnect_skips_within_interval(opc_repository):
|
||||
assert opc_repository._reconnect_task is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_reconnect_skips_when_already_in_progress(opc_repository):
|
||||
opc_repository.last_reconnection_time = None
|
||||
in_progress_task = asyncio.create_task(asyncio.sleep(60))
|
||||
opc_repository._reconnect_task = in_progress_task
|
||||
|
||||
await opc_repository._start_reconnect('BadSessionIdInvalid', 'tok')
|
||||
|
||||
assert opc_repository._reconnect_task is in_progress_task
|
||||
opc_repository.warning.assert_called_once()
|
||||
|
||||
in_progress_task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await in_progress_task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_data_protocol_closed_schedules_reconnect(opc_repository):
|
||||
opc_repository.client = MagicMock()
|
||||
@@ -608,3 +778,23 @@ async def test_reconnect_locked_sets_last_reconnection_time(mock_datetime, opc_r
|
||||
opc_repository._connect_locked.assert_called_once()
|
||||
assert result == (True, {})
|
||||
assert opc_repository.last_reconnection_time == datetime(2025, 1, 1, 12, 0, 0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_reconnect_failed_without_error_detail_logs_once(opc_repository):
|
||||
"""A falsy `error` (no message to add) must not log a second, empty line."""
|
||||
opc_repository._reconnect_locked = AsyncMock(return_value=(False, {}))
|
||||
|
||||
await opc_repository._run_reconnect('BadSessionIdInvalid', 'tok')
|
||||
|
||||
assert opc_repository.error.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_reconnect_swallows_an_unexpected_exception(opc_repository):
|
||||
"""A bug inside the reconnect itself must not escape the background task."""
|
||||
opc_repository._reconnect_locked = AsyncMock(side_effect=RuntimeError('boom'))
|
||||
|
||||
await opc_repository._run_reconnect('BadSessionIdInvalid', 'tok')
|
||||
|
||||
assert opc_repository.error.call_count >= 2
|
||||
|
||||
40
tests/laborious/utils/test_dataframe_debug.py
Normal file
40
tests/laborious/utils/test_dataframe_debug.py
Normal file
@@ -0,0 +1,40 @@
|
||||
"""Debug message formatting: full dataframe when small, a summary when it would flood the log."""
|
||||
|
||||
from pandas import DataFrame
|
||||
|
||||
from laborious.utils.dataframe_debug import (
|
||||
DEFAULT_MAX_DEBUG_DATAFRAME_ROWS,
|
||||
build_dataframe_debug_message,
|
||||
)
|
||||
|
||||
|
||||
def test_non_dataframe_payloads_are_stringified():
|
||||
assert build_dataframe_debug_message('payload:', {'a': 1}) == "payload: {'a': 1}"
|
||||
|
||||
|
||||
def test_a_small_dataframe_is_logged_in_full():
|
||||
data = DataFrame({'a': [1, 2]})
|
||||
|
||||
message = build_dataframe_debug_message('payload:', data)
|
||||
|
||||
assert message.startswith('payload:\n')
|
||||
assert data.to_csv() in message
|
||||
|
||||
|
||||
def test_a_dataframe_over_the_row_limit_is_summarized_instead():
|
||||
data = DataFrame({'a': range(DEFAULT_MAX_DEBUG_DATAFRAME_ROWS + 1)})
|
||||
|
||||
message = build_dataframe_debug_message('payload:', data)
|
||||
|
||||
assert 'skipped because dataframe has' in message
|
||||
assert str(data.shape) in message
|
||||
assert data.to_csv() not in message
|
||||
|
||||
|
||||
def test_the_row_limit_is_configurable():
|
||||
data = DataFrame({'a': [1, 2, 3]})
|
||||
|
||||
message = build_dataframe_debug_message('payload:', data, max_rows=1)
|
||||
|
||||
assert 'skipped because dataframe has 3 rows' in message
|
||||
assert '(max: 1)' in message
|
||||
@@ -1,8 +1,11 @@
|
||||
"""The password envelope: AES-256-GCM in, plaintext password out, nothing in between."""
|
||||
|
||||
import base64
|
||||
import os
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
||||
|
||||
from laborious.utils.import_password import (
|
||||
EnvelopeError,
|
||||
@@ -89,6 +92,26 @@ def test_rejects_an_unusable_key(raw_key, match):
|
||||
load_envelope_key(raw_key)
|
||||
|
||||
|
||||
@patch('laborious.utils.import_password.AESGCM')
|
||||
def test_a_generic_decrypt_failure_is_reported_as_could_not_be_opened(mock_aesgcm):
|
||||
"""Anything AESGCM raises besides InvalidTag (e.g. a malformed nonce) is not authentication."""
|
||||
mock_aesgcm.return_value.decrypt.side_effect = ValueError('bad nonce length')
|
||||
envelope = factory.build_envelope('Correct1Horse!Battery', KEY)
|
||||
|
||||
with pytest.raises(EnvelopeError, match='could not be opened'):
|
||||
decrypt_password_envelope(envelope, KEY)
|
||||
|
||||
|
||||
def test_a_non_utf8_plaintext_is_rejected():
|
||||
"""The envelope opened but what came out is not text — refuse rather than mangle it."""
|
||||
nonce = os.urandom(12)
|
||||
sealed = AESGCM(KEY).encrypt(nonce, b'\xff\xfe\xfd', None)
|
||||
envelope = base64.b64encode(nonce + sealed).decode('ascii')
|
||||
|
||||
with pytest.raises(EnvelopeError, match='not valid UTF-8'):
|
||||
decrypt_password_envelope(envelope, KEY)
|
||||
|
||||
|
||||
def test_no_secret_appears_in_the_error_text():
|
||||
password = 'Correct1Horse!Battery'
|
||||
envelope = factory.build_envelope(password, KEY)
|
||||
|
||||
245
tests/laborious/worker/test_worker_main.py
Normal file
245
tests/laborious/worker/test_worker_main.py
Normal file
@@ -0,0 +1,245 @@
|
||||
"""``main()`` and ``start_prometheus_server()``: the orchestration paths ``test_runtime_task_queues``
|
||||
doesn't reach.
|
||||
|
||||
Everything I/O-bound (Temporal client, Prometheus HTTP server, Mongo/Postgres/MinIO configs) is
|
||||
mocked; only the control flow inside ``worker.py`` is under test.
|
||||
"""
|
||||
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
||||
|
||||
from pytest import mark, raises
|
||||
|
||||
from laborious.worker.worker import main, start_prometheus_server
|
||||
|
||||
|
||||
@patch('laborious.worker.worker.metrics.APP_UP')
|
||||
@patch('laborious.worker.worker.get_logger')
|
||||
def test_main_exits_when_runtime_is_unset(mock_get_logger, mock_app_up, monkeypatch):
|
||||
monkeypatch.delenv('RUNTIME', raising=False)
|
||||
logger = MagicMock()
|
||||
mock_get_logger.return_value = logger
|
||||
|
||||
with raises(SystemExit) as exc:
|
||||
import asyncio
|
||||
|
||||
asyncio.run(main())
|
||||
|
||||
assert exc.value.code == 1
|
||||
logger.custom_critical.assert_called_once()
|
||||
mock_app_up.labels.return_value.set.assert_called_once_with(0)
|
||||
|
||||
|
||||
@patch('laborious.worker.worker.prepare_worker')
|
||||
@patch('laborious.worker.worker.client.Client.connect', new_callable=AsyncMock)
|
||||
@patch('laborious.worker.worker.TelemetryConfig')
|
||||
@patch('laborious.worker.worker.PrometheusConfig')
|
||||
@patch('laborious.worker.worker.Runtime')
|
||||
@patch('laborious.worker.worker.Activities')
|
||||
@patch('laborious.worker.worker.NotificationHandler')
|
||||
@patch('laborious.worker.worker.build_import_status_config')
|
||||
@patch('laborious.worker.worker.build_import_config')
|
||||
@patch('laborious.worker.worker.build_api_config')
|
||||
@patch('laborious.worker.worker.build_opc_config')
|
||||
@patch('laborious.worker.worker.build_minio_config')
|
||||
@patch('laborious.worker.worker.build_mlflow_config')
|
||||
@patch('laborious.worker.worker.build_postgres_config')
|
||||
@patch('laborious.worker.worker.build_mongodb_config')
|
||||
@patch('laborious.worker.worker.start_prometheus_server')
|
||||
@patch('laborious.worker.worker.get_logger')
|
||||
class TestMainHappyPath:
|
||||
"""One class so every test shares the same (long) patch stack."""
|
||||
|
||||
def _run(self, mock_prepare_worker, mock_connect):
|
||||
worker = MagicMock()
|
||||
worker.run = AsyncMock(return_value=None)
|
||||
mock_prepare_worker.return_value = worker
|
||||
mock_connect.return_value = MagicMock()
|
||||
|
||||
import asyncio
|
||||
|
||||
with raises(SystemExit) as exc:
|
||||
asyncio.run(main())
|
||||
return exc.value.code
|
||||
|
||||
def test_runs_four_workers_without_the_import_flag( # noqa: PLR0913
|
||||
self,
|
||||
mock_get_logger,
|
||||
_mock_start_prometheus,
|
||||
_mock_mongo_cfg,
|
||||
_mock_pg_cfg,
|
||||
_mock_mlflow_cfg,
|
||||
_mock_minio_cfg,
|
||||
_mock_opc_cfg,
|
||||
_mock_api_cfg,
|
||||
_mock_import_cfg,
|
||||
_mock_import_status_cfg,
|
||||
mock_notification_handler,
|
||||
mock_activities,
|
||||
_mock_runtime,
|
||||
_mock_prometheus_config,
|
||||
_mock_telemetry_config,
|
||||
mock_connect,
|
||||
mock_prepare_worker,
|
||||
monkeypatch,
|
||||
):
|
||||
monkeypatch.setenv('RUNTIME', 'test-runtime')
|
||||
monkeypatch.delenv('IMPORT_MODEL_ENABLED', raising=False)
|
||||
mock_get_logger.return_value = MagicMock()
|
||||
mock_activities.return_value.init_opc = AsyncMock()
|
||||
mock_activities.return_value.shutdown = AsyncMock()
|
||||
|
||||
exit_code = self._run(mock_prepare_worker, mock_connect)
|
||||
|
||||
assert exit_code == 0
|
||||
assert mock_prepare_worker.call_count == 4
|
||||
mock_activities.return_value.init_opc.assert_awaited_once()
|
||||
mock_activities.return_value.shutdown.assert_awaited_once()
|
||||
mock_notification_handler.return_value.shutdown.assert_called_once()
|
||||
|
||||
def test_runs_five_workers_with_the_import_flag( # noqa: PLR0913
|
||||
self,
|
||||
mock_get_logger,
|
||||
_mock_start_prometheus,
|
||||
_mock_mongo_cfg,
|
||||
_mock_pg_cfg,
|
||||
_mock_mlflow_cfg,
|
||||
_mock_minio_cfg,
|
||||
_mock_opc_cfg,
|
||||
_mock_api_cfg,
|
||||
_mock_import_cfg,
|
||||
_mock_import_status_cfg,
|
||||
_mock_notification_handler,
|
||||
mock_activities,
|
||||
_mock_runtime,
|
||||
_mock_prometheus_config,
|
||||
_mock_telemetry_config,
|
||||
mock_connect,
|
||||
mock_prepare_worker,
|
||||
monkeypatch,
|
||||
):
|
||||
monkeypatch.setenv('RUNTIME', 'test-runtime')
|
||||
monkeypatch.setenv('IMPORT_MODEL_ENABLED', 'true')
|
||||
mock_get_logger.return_value = MagicMock()
|
||||
mock_activities.return_value.init_opc = AsyncMock()
|
||||
mock_activities.return_value.shutdown = AsyncMock()
|
||||
|
||||
exit_code = self._run(mock_prepare_worker, mock_connect)
|
||||
|
||||
assert exit_code == 0
|
||||
assert mock_prepare_worker.call_count == 5
|
||||
|
||||
def test_a_falsy_notification_handler_or_activities_skips_shutdown( # noqa: PLR0913
|
||||
self,
|
||||
mock_get_logger,
|
||||
_mock_start_prometheus,
|
||||
_mock_mongo_cfg,
|
||||
_mock_pg_cfg,
|
||||
_mock_mlflow_cfg,
|
||||
_mock_minio_cfg,
|
||||
_mock_opc_cfg,
|
||||
_mock_api_cfg,
|
||||
_mock_import_cfg,
|
||||
_mock_import_status_cfg,
|
||||
mock_notification_handler,
|
||||
mock_activities,
|
||||
_mock_runtime,
|
||||
_mock_prometheus_config,
|
||||
_mock_telemetry_config,
|
||||
mock_connect,
|
||||
mock_prepare_worker,
|
||||
monkeypatch,
|
||||
):
|
||||
"""Defensive `if notification_handler:` / `if activities:` guards around shutdown — real
|
||||
construction always yields truthy instances, but the guard itself still needs a test."""
|
||||
monkeypatch.setenv('RUNTIME', 'test-runtime')
|
||||
monkeypatch.delenv('IMPORT_MODEL_ENABLED', raising=False)
|
||||
mock_get_logger.return_value = MagicMock()
|
||||
mock_activities.return_value.init_opc = AsyncMock()
|
||||
mock_activities.return_value.shutdown = AsyncMock()
|
||||
mock_activities.return_value.__bool__.return_value = False
|
||||
mock_notification_handler.return_value.__bool__.return_value = False
|
||||
|
||||
exit_code = self._run(mock_prepare_worker, mock_connect)
|
||||
|
||||
assert exit_code == 0
|
||||
mock_activities.return_value.shutdown.assert_not_called()
|
||||
mock_notification_handler.return_value.shutdown.assert_not_called()
|
||||
|
||||
def test_a_worker_failure_exits_non_zero_and_still_cleans_up( # noqa: PLR0913
|
||||
self,
|
||||
mock_get_logger,
|
||||
_mock_start_prometheus,
|
||||
_mock_mongo_cfg,
|
||||
_mock_pg_cfg,
|
||||
_mock_mlflow_cfg,
|
||||
_mock_minio_cfg,
|
||||
_mock_opc_cfg,
|
||||
_mock_api_cfg,
|
||||
_mock_import_cfg,
|
||||
_mock_import_status_cfg,
|
||||
mock_notification_handler,
|
||||
mock_activities,
|
||||
_mock_runtime,
|
||||
_mock_prometheus_config,
|
||||
_mock_telemetry_config,
|
||||
mock_connect,
|
||||
mock_prepare_worker,
|
||||
monkeypatch,
|
||||
):
|
||||
monkeypatch.setenv('RUNTIME', 'test-runtime')
|
||||
monkeypatch.delenv('IMPORT_MODEL_ENABLED', raising=False)
|
||||
logger = MagicMock()
|
||||
mock_get_logger.return_value = logger
|
||||
mock_activities.return_value.init_opc = AsyncMock()
|
||||
mock_activities.return_value.shutdown = AsyncMock()
|
||||
mock_connect.return_value = MagicMock()
|
||||
|
||||
worker = MagicMock()
|
||||
worker.run = AsyncMock(side_effect=RuntimeError('boom'))
|
||||
mock_prepare_worker.return_value = worker
|
||||
|
||||
import asyncio
|
||||
|
||||
with raises(SystemExit) as exc:
|
||||
asyncio.run(main())
|
||||
|
||||
assert exc.value.code == 1
|
||||
logger.custom_error.assert_called_once_with(ANY, ANY)
|
||||
mock_activities.return_value.shutdown.assert_awaited_once()
|
||||
mock_notification_handler.return_value.shutdown.assert_called_once()
|
||||
|
||||
|
||||
@patch('laborious.worker.worker.metrics.APP_UP')
|
||||
@patch('laborious.worker.worker.start_http_server')
|
||||
def test_start_prometheus_server_marks_the_app_up(mock_start_http_server, mock_app_up):
|
||||
start_prometheus_server()
|
||||
|
||||
mock_start_http_server.assert_called_once()
|
||||
mock_app_up.labels.return_value.set.assert_called_once_with(1)
|
||||
|
||||
|
||||
@patch('laborious.worker.worker.os._exit')
|
||||
@patch('laborious.worker.worker.metrics.APP_UP')
|
||||
@patch('laborious.worker.worker.start_http_server')
|
||||
def test_start_prometheus_server_exits_the_process_on_failure(
|
||||
mock_start_http_server, mock_app_up, mock_os_exit
|
||||
):
|
||||
mock_start_http_server.side_effect = OSError('port in use')
|
||||
|
||||
start_prometheus_server()
|
||||
|
||||
mock_os_exit.assert_called_once_with(1)
|
||||
mock_app_up.labels.return_value.set.assert_not_called()
|
||||
|
||||
|
||||
@mark.parametrize('port_env', ['9999'])
|
||||
@patch('laborious.worker.worker.metrics.APP_UP')
|
||||
@patch('laborious.worker.worker.start_http_server')
|
||||
def test_start_prometheus_server_reads_the_configured_port(
|
||||
mock_start_http_server, _mock_app_up, port_env, monkeypatch
|
||||
):
|
||||
monkeypatch.setenv('HTTP_METRICS_PORT', port_env)
|
||||
|
||||
start_prometheus_server()
|
||||
|
||||
mock_start_http_server.assert_called_once_with(int(port_env))
|
||||
@@ -1,6 +1,6 @@
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
|
||||
|
||||
from pytest import fixture, mark
|
||||
from pytest import fixture, mark, raises
|
||||
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
|
||||
@@ -791,6 +791,30 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
|
||||
workflow_mock.execute_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_run_cleans_up_and_reraises_on_pipeline_failure(workflow_mock, prediction_process):
|
||||
"""A failure anywhere in the pipeline must still release the MinIO offload before the
|
||||
workflow fails — cleanup is not conditional on success."""
|
||||
prediction_process._run_prediction_pipeline = AsyncMock(side_effect=RuntimeError('boom'))
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'data': {'some': 'data'},
|
||||
'model_id': 1,
|
||||
'model_name': 'test_model_name',
|
||||
}
|
||||
|
||||
with raises(RuntimeError, match='boom'):
|
||||
await prediction_process.run(input_data)
|
||||
|
||||
workflow_mock.execute_activity_method.assert_called_once_with(
|
||||
Activities.cleanup_minio_objects_expired,
|
||||
{**metadata, 'data': input_data['data']},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_run_with_cleanup_prefixes(workflow_mock, prediction_process):
|
||||
|
||||
Reference in New Issue
Block a user