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.

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()
 
 

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):

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',

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

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

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

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
 
 

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
 
 

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',

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

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

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)

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))

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):