diff --git a/pyproject.toml b/pyproject.toml index 2f2ef0d..7ffcad0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -179,6 +179,7 @@ exclude_lines = [ "if TYPE_CHECKING:", "class .*\\bProtocol\\):", "@(abc\\.)?abstractmethod", + "@overload", ] [tool.coverage.html] diff --git a/tests/helpers/bundle_factory.py b/tests/helpers/bundle_factory.py index db54be5..d7ce3e3 100644 --- a/tests/helpers/bundle_factory.py +++ b/tests/helpers/bundle_factory.py @@ -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() diff --git a/tests/laborious/activities/test_api.py b/tests/laborious/activities/test_api.py index bf5b9f1..45d5e92 100644 --- a/tests/laborious/activities/test_api.py +++ b/tests/laborious/activities/test_api.py @@ -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): diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index 5d242e5..da17c94 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -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', diff --git a/tests/laborious/activities/test_model_import_provisioning.py b/tests/laborious/activities/test_model_import_provisioning.py index 7915a04..fe5d11f 100644 --- a/tests/laborious/activities/test_model_import_provisioning.py +++ b/tests/laborious/activities/test_model_import_provisioning.py @@ -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 diff --git a/tests/laborious/activities/test_model_import_status.py b/tests/laborious/activities/test_model_import_status.py index 998d855..2f24080 100644 --- a/tests/laborious/activities/test_model_import_status.py +++ b/tests/laborious/activities/test_model_import_status.py @@ -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 diff --git a/tests/laborious/activities/test_opc.py b/tests/laborious/activities/test_opc.py index 6a36f8a..2e248a6 100644 --- a/tests/laborious/activities/test_opc.py +++ b/tests/laborious/activities/test_opc.py @@ -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 diff --git a/tests/laborious/utils/bundle/test_reader.py b/tests/laborious/utils/bundle/test_reader.py index 81dd2bd..ebff1bf 100644 --- a/tests/laborious/utils/bundle/test_reader.py +++ b/tests/laborious/utils/bundle/test_reader.py @@ -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 diff --git a/tests/laborious/utils/models/test_import_errors.py b/tests/laborious/utils/models/test_import_errors.py index ee949b7..53472fa 100644 --- a/tests/laborious/utils/models/test_import_errors.py +++ b/tests/laborious/utils/models/test_import_errors.py @@ -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 diff --git a/tests/laborious/utils/models/test_minio_dataframe_payload.py b/tests/laborious/utils/models/test_minio_dataframe_payload.py index c27765d..88d7885 100644 --- a/tests/laborious/utils/models/test_minio_dataframe_payload.py +++ b/tests/laborious/utils/models/test_minio_dataframe_payload.py @@ -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', diff --git a/tests/laborious/utils/repository/test_opc_repository.py b/tests/laborious/utils/repository/test_opc_repository.py index 365a400..f9d489b 100644 --- a/tests/laborious/utils/repository/test_opc_repository.py +++ b/tests/laborious/utils/repository/test_opc_repository.py @@ -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 diff --git a/tests/laborious/utils/test_dataframe_debug.py b/tests/laborious/utils/test_dataframe_debug.py new file mode 100644 index 0000000..0d0d6c5 --- /dev/null +++ b/tests/laborious/utils/test_dataframe_debug.py @@ -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 diff --git a/tests/laborious/utils/test_import_password.py b/tests/laborious/utils/test_import_password.py index a852861..eaf7a5a 100644 --- a/tests/laborious/utils/test_import_password.py +++ b/tests/laborious/utils/test_import_password.py @@ -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) diff --git a/tests/laborious/worker/test_worker_main.py b/tests/laborious/worker/test_worker_main.py new file mode 100644 index 0000000..e87aac7 --- /dev/null +++ b/tests/laborious/worker/test_worker_main.py @@ -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)) diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py index a018504..cdef6cd 100644 --- a/tests/laborious/workflows/subworkflows/test_prediction_process.py +++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py @@ -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):