sientia-dataops-laborious-legacy/e2e/test_import_model_provisioning_failures.py

"""Each provisioning step fails as itself: its own code, its own sentence, nothing half-built.

The four steps after the gates — experiment creation, artifact upload, registration, the model
document — each get one failure, and every one lands `IMPORT_ERROR` rather than
`ORCHESTRATOR_VALIDATION_ERROR`. That split is not cosmetic: it is the table's existing distinction
between "the person who submitted this can fix it" and "only support can", and a provisioning
failure is always the second kind.

**How each failure is caused**, in design E7's order of preference:

| Step | Cause | Preference |
|---|---|---|
| `experiment_creation` | the MLflow server is genuinely stopped | broken dependency |
| `artifact_upload` | targeted patch on `log_artifacts` | patch, justified in place |
| `registration` | targeted patch on `create_model_version` | patch, justified in place |
| `model_document` | a real unique index refuses the activity's own insert | broken dependency |

Two departures from what the tasks anticipated, both recorded in the test that found them:

- **Registration needs a patch too**, not just artifact upload. One process serves both tracking and
  the registry, so stopping it fails experiment creation three steps earlier; and MLflow validates
  nothing about a registered model name except that it is non-empty, which gate 6 already requires.
  Both patches are on the single client call the step *is*, never on the activity or the workflow, so
  what is under test — the workflow attributing this step's failure to this step — is untouched.
- **The model document's unique-index violation is on `model_config.target`**, not on `id`. The
  duplicate-`id` state the task describes cannot be constructed at all; `TestModelDocument`'s
  docstring shows why, and the mechanism the task wanted — a real index refusing a real insert — is
  kept intact.
"""

from __future__ import annotations

import json
import tempfile
from typing import Any

import pytest
from temporalio.client import WorkflowFailureError

from e2e.import_harness import IMPORT_BUCKET, IMPORT_PREFIX, SEEDED_MODEL_IDS
from laborious.activities.model_import_errors import IMPORT_FAILURE_CATALOG
from laborious.utils.bundle.steps import ImportStep

pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]

PLATFORM_ERROR_STATUS = 'IMPORT_ERROR'
FACTORY_PASSWORD = 'Factory1Model!Test'


def _as_dict(value: Any) -> dict[str, Any]:
    return json.loads(value) if isinstance(value, str) else (value or {})


def _assert_platform_failure(
    *,
    row: dict[str, Any] | None,
    step: ImportStep,
    workflow_id: str,
    object_key: str,
) -> None:
    """The shared claim of task 7.5, asserted by every case in this module.

    `error_message` is the catalog sentence *exactly* — not "contains" — which is the strongest
    possible form of "carries no path, table name, object key or traceback": there is nothing in it
    that did not come from the catalog.
    """
    assert row is not None, 'the failure lost the row it was recorded on'
    assert row['status'] == PLATFORM_ERROR_STATUS, (
        f"status is {row['status']!r}; a provisioning failure is not the user's to fix"
    )

    code, reason = IMPORT_FAILURE_CATALOG[step]
    detail = _as_dict(row['orchestrator_response_data'])['import']
    assert detail['step'] == step.value, detail
    assert detail['code'] == code.value, detail
    assert detail['workflow_id'] == workflow_id, detail
    assert row['error_message'] == reason

    message = row['error_message'] or ''
    for forbidden in (
        object_key,
        IMPORT_BUCKET,
        IMPORT_PREFIX,
        'experiment_run',
        'Traceback',
        'Error',
        'mlflow',
        'Mongo',
        tempfile.gettempdir(),
    ):
        assert forbidden not in message, f'{forbidden!r} leaked into error_message: {message!r}'


async def _upload_and_start(
    *,
    import_bundle_factory: Any,
    import_upload_bundle: Any,
    import_insert_status_row: Any,
    import_build_input: Any,
    import_execute: Any,
    import_envelope: Any,
    **factory_kwargs: Any,
) -> tuple[Any, int, str]:
    object_key, digest = import_upload_bundle(import_bundle_factory(**factory_kwargs))
    import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
    handle = await import_execute(
        import_build_input(
            import_run_id=import_run_id,
            object_key=object_key,
            expected_digest=digest,
            password_envelope=import_envelope(FACTORY_PASSWORD),
        )
    )
    return handle, import_run_id, object_key


class TestExperimentCreation:
    """The tracking server is not there. No patch, no mock — the process is stopped."""

    async def test_a_stopped_tracking_server_fails_at_experiment_creation(
        self,
        import_status_table: Any,
        import_models_collection: Any,
        import_names: dict[str, str],
        import_bundle_factory: Any,
        import_upload_bundle: Any,
        import_insert_status_row: Any,
        import_build_input: Any,
        import_execute: Any,
        import_envelope: Any,
        import_read_status_row: Any,
        import_mlflow_stopped: Any,
        import_assert_nothing_provisioned: Any,
    ) -> None:
        with import_mlflow_stopped():
            handle, import_run_id, object_key = await _upload_and_start(
                import_bundle_factory=import_bundle_factory,
                import_upload_bundle=import_upload_bundle,
                import_insert_status_row=import_insert_status_row,
                import_build_input=import_build_input,
                import_execute=import_execute,
                import_envelope=import_envelope,
            )
            with pytest.raises(WorkflowFailureError):
                await handle.result()
            row = import_read_status_row(import_run_id)

        _assert_platform_failure(
            row=row,
            step=ImportStep.EXPERIMENT_CREATION,
            workflow_id=handle.id,
            object_key=object_key,
        )

        # The bundle *was* opened before MLflow was reached, so the names are on the row — which is
        # what task 9.3 asks for, observed from a failure that stopped exactly there.
        assert row is not None
        assert row['experiment_name'] == import_names['experiment_name']
        assert row['run_name'] == f'import-{import_names["model_name"]}-v2'

        # Nothing in Mongo: the document is the last step and never ran.
        documents = list(import_models_collection.find({}, {'_id': False}))
        assert {str(d['id']) for d in documents} == set(SEEDED_MODEL_IDS)
        import_assert_nothing_provisioned(
            experiment_name=import_names['experiment_name'],
            model_name=import_names['model_name'],
            object_key=object_key,
        )


class TestArtifactUpload:
    """The step design E7 names as the one a patch is needed for."""

    async def test_a_failing_artifact_upload_is_recorded_as_its_own_step(
        self,
        monkeypatch: pytest.MonkeyPatch,
        import_status_table: Any,
        import_models_collection: Any,
        import_names: dict[str, str],
        import_bundle_factory: Any,
        import_activities: Any,
        import_mlflow_server: dict[str, Any],
        import_upload_bundle: Any,
        import_insert_status_row: Any,
        import_build_input: Any,
        import_execute: Any,
        import_envelope: Any,
        import_read_status_row: Any,
    ) -> None:
        """Why a patch: with the server up and the bundle valid, nothing else fails only this step.

        Hostile input cannot reach here — a bundle that gets past gate 7 has a readable tree, and
        one that does not never reaches artifact upload. A broken dependency cannot isolate it
        either: the only dependency this step has is the tracking server, and stopping that fails
        experiment creation one step earlier. So the patch is on the single client call the step
        makes, `log_artifacts`; everything around it — the run creation, the parameters, the
        workflow's attribution of the failure — is the real code.
        """
        client = import_activities.import_mlflow_repository.client

        def refuse(*_args: Any, **_kwargs: Any) -> None:
            raise OSError('artifact store is unreachable')

        monkeypatch.setattr(client, 'log_artifacts', refuse)

        handle, import_run_id, object_key = await _upload_and_start(
            import_bundle_factory=import_bundle_factory,
            import_upload_bundle=import_upload_bundle,
            import_insert_status_row=import_insert_status_row,
            import_build_input=import_build_input,
            import_execute=import_execute,
            import_envelope=import_envelope,
        )
        with pytest.raises(WorkflowFailureError):
            await handle.result()

        _assert_platform_failure(
            row=import_read_status_row(import_run_id),
            step=ImportStep.ARTIFACT_UPLOAD,
            workflow_id=handle.id,
            object_key=object_key,
        )

        # The experiment exists — it was the previous step and it succeeded. Nothing after does.
        mlflow_client = _fresh_client(import_mlflow_server)
        assert mlflow_client.get_experiment_by_name(import_names['experiment_name']) is not None
        assert not list(mlflow_client.search_model_versions(f"name='{import_names['model_name']}'"))
        assert import_models_collection.find_one({'name': import_names['model_name']}) is None


class TestRegistration:
    """The artifacts are stored and the model still cannot be registered."""

    async def test_a_failing_registration_is_recorded_as_its_own_step(
        self,
        monkeypatch: pytest.MonkeyPatch,
        import_status_table: Any,
        import_models_collection: Any,
        import_names: dict[str, str],
        import_bundle_factory: Any,
        import_activities: Any,
        import_mlflow_server: dict[str, Any],
        import_upload_bundle: Any,
        import_insert_status_row: Any,
        import_build_input: Any,
        import_execute: Any,
        import_envelope: Any,
        import_read_status_row: Any,
    ) -> None:
        """Why a patch here too, contrary to design E7's expectation that this step needs none.

        Hostile input was the first thing tried: a registered model name MLflow refuses. It does not
        exist — `_validate_model_name` rejects only the empty string, which gate 6 already refuses,
        and the sqlite backend store does not enforce the `VARCHAR(256)` the schema declares, so
        even a 300-character name is accepted. A broken dependency was the second: one process
        serves both tracking and the registry, so stopping it fails experiment creation three steps
        earlier. The patch is on `create_model_version`, the one call this step is.
        """
        client = import_activities.import_mlflow_repository.client

        def refuse(*_args: Any, **_kwargs: Any) -> None:
            raise RuntimeError('registry write rejected')

        monkeypatch.setattr(client, 'create_model_version', refuse)

        handle, import_run_id, object_key = await _upload_and_start(
            import_bundle_factory=import_bundle_factory,
            import_upload_bundle=import_upload_bundle,
            import_insert_status_row=import_insert_status_row,
            import_build_input=import_build_input,
            import_execute=import_execute,
            import_envelope=import_envelope,
        )
        with pytest.raises(WorkflowFailureError):
            await handle.result()

        _assert_platform_failure(
            row=import_read_status_row(import_run_id),
            step=ImportStep.REGISTRATION,
            workflow_id=handle.id,
            object_key=object_key,
        )

        # The run and its artifacts survive — "the model's files were stored, but the model itself
        # could not be registered" is exactly what the catalog sentence for this step says.
        mlflow_client = _fresh_client(import_mlflow_server)
        experiment = mlflow_client.get_experiment_by_name(import_names['experiment_name'])
        assert experiment is not None
        runs = mlflow_client.search_runs([experiment.experiment_id])
        assert len(runs) == 1
        assert {a.path for a in mlflow_client.list_artifacts(runs[0].info.run_id)} >= {
            'prediction_model',
            'data_model',
        }
        assert import_models_collection.find_one({'name': import_names['model_name']}) is None


class TestModelDocument:
    """The last step, failed by a real unique index refusing the real insert.

    **Why the index is on `model_config.target` and not on `id`.** Task 7.4 asks for "a document
    already holding the id the activity is about to compute". That state cannot be constructed, and
    why is worth more than the test would have been: `_next_model_document_id` computes
    `max(int(id)) + 1` over `distinct('id')` of the *same* collection, so every value that
    stringifies to the colliding id is counted by that maximum. Put a document at `'8'` and the next
    id becomes `'9'`. Make the id unparseable and it is ignored by the numbering *and* does not
    collide. There is no third case — the numbering cannot collide with itself, and the only real
    duplicate-key path on `id` is two imports racing, which this change lists as a non-goal. That is
    a good property of the implementation, not a gap.

    So the mechanism the task asked for is kept — a real unique index refusing the activity's own
    insert — on a key that can actually collide. `model_config.target` is `'SE'` on the seeded
    documents and `'SE'` on the one this import writes, so a unique index there makes the insert a
    genuine `DuplicateKeyError` from the server. No patch, no mock: the same class of failure, from
    the same enforcement mechanism, at the same step.
    """

    async def test_a_unique_index_violation_fails_at_model_document(
        self,
        import_status_table: Any,
        import_models_collection: Any,
        import_names: dict[str, str],
        import_bundle_factory: Any,
        import_mlflow_server: dict[str, Any],
        import_upload_bundle: Any,
        import_insert_status_row: Any,
        import_build_input: Any,
        import_execute: Any,
        import_envelope: Any,
        import_read_status_row: Any,
    ) -> None:
        # All three seeded documents carry target 'SE', so the index has to be buildable first:
        # give two of them their own target and leave exactly one holding 'SE'. That one is what the
        # import's own document — also target 'SE' — collides with.
        for ordinal, seeded_id in enumerate(sorted(SEEDED_MODEL_IDS)[1:], start=1):
            import_models_collection.update_one(
                {'id': seeded_id}, {'$set': {'model_config.target': f'OTHER_{ordinal}'}}
            )
        import_models_collection.create_index('model_config.target', unique=True, name='e2e_target')
        try:
            handle, import_run_id, object_key = await _upload_and_start(
                import_bundle_factory=import_bundle_factory,
                import_upload_bundle=import_upload_bundle,
                import_insert_status_row=import_insert_status_row,
                import_build_input=import_build_input,
                import_execute=import_execute,
                import_envelope=import_envelope,
                parameters={'target_variable': 'SE'},
            )
            with pytest.raises(WorkflowFailureError):
                await handle.result()
            row = import_read_status_row(import_run_id)
        finally:
            # The collection is dropped and reseeded per test, but dropping the index here keeps the
            # failure local even if that ever changes.
            import_models_collection.drop_index('e2e_target')

        _assert_platform_failure(
            row=row,
            step=ImportStep.MODEL_DOCUMENT,
            workflow_id=handle.id,
            object_key=object_key,
        )

        # No document for this model, and the seeded ones are exactly as they were.
        assert import_models_collection.find_one({'name': import_names['model_name']}) is None
        assert {
            str(document['id']) for document in import_models_collection.find({}, {'_id': False})
        } == set(SEEDED_MODEL_IDS)

        # Everything before it stands — registered, and unusable, which is the point of the split
        # between this step's sentence and the registration step's.
        mlflow_client = _fresh_client(import_mlflow_server)
        versions = list(mlflow_client.search_model_versions(f"name='{import_names['model_name']}'"))
        assert len(versions) == 1


def _fresh_client(server: dict[str, Any]) -> Any:
    """An MlflowClient bound to the test server, independent of the one the activities hold."""
    import mlflow

    return mlflow.tracking.MlflowClient(tracking_uri=server['tracking_uri'])