"""Fixtures for the `.sientia` model-import e2e suite. Registered from `e2e/conftest.py`. **Why this is a module of its own.** `e2e/conftest.py` is shared by every existing e2e suite and long enough already; four containers' worth of setup, an MLflow server subprocess, the import environment variables and an `ImportModel` worker do not belong in it. Every fixture here is prefixed `import_`, so nothing collides with the predictions harness, and no existing fixture is edited — the one exception being `log_statement=all` on the shared `postgres_container`, which the absence assertions need and which changes no test's behaviour. **The fidelity contract, which is the whole point of this file.** On the import path *nothing is mocked*: | Dependency | Here | Why real | |---|---|---| | MinIO | the shared `minio_container`, plus an import bucket | gate 1 reads a real object | | BFF Postgres | a container of its own, `sientia-core-mlops-bff` | the real `CHECK` constraints are half of what is under test, and two tests stop the server | | MongoDB | a container, `models` seeded with gapped ids and a unique index on `id` | `max(int(id)) + 1` against a real unique index | | MLflow | `mlflow server` from *this* venv, sqlite-backed, artifacts proxied | experiment/run/artifact/registry semantics are what the provisioning activities are tested for | | the crypto | the production reader, real Argon2id, real `secretstream` | a bundle the producer wrote must open | What stays stubbed, deliberately: `CoreNotificationHandler`'s Mongo client keeps the existing `mock_mongo_client`. Notifications are a side channel this suite counts, not a dependency it tests, and a real notification collection would add a second write path to every absence assertion for nothing. The import `Activities` is **never** composed with `patch_mlflow`, `patch_minio_repository` or `patch_create_engine`. `mlflow.set_tracking_uri` is process-global, so a leak would silently turn a real assertion into an assertion about a `MagicMock`; `import_activities` asserts the live tracking URI names the test server before it hands the object over. **Bundles come from the unit suite.** `tests/helpers/bundle_factory.py` writes them and `tests/fixtures/bundle/golden_model_v3.sientia` is the producer-emitted one. No third mirrored copy of the producer's crypto is introduced here. **Isolation** (the containers are session-scoped): per test the status table is rebuilt, the `models` collection is dropped and reseeded, and the import prefix is emptied. MLflow state is *not* reset — experiment and model names are made unique per test instead, because deleted registered models leave tombstones `get_registered_model` still finds, which would make "second import of the same model" unreadable. """ from __future__ import annotations import contextlib import hashlib import json import os import signal import socket import subprocess import sys import time import uuid from collections.abc import Callable, Iterator from datetime import UTC, datetime from pathlib import Path from typing import Any import mlflow import pytest import pytest_asyncio import requests from pymongo import MongoClient from sqlalchemy import Engine, create_engine, text from temporalio.client import WorkflowHandle from temporalio.worker import Worker from testcontainers.mongodb import MongoDbContainer from testcontainers.postgres import PostgresContainer from laborious.activities.activities import Activities from laborious.workflows.import_model import ImportModel from tests.helpers.bundle_factory import ( DEFAULT_ENVELOPE_KEY, DEFAULT_PASSWORD, ZipSpec, build_envelope, build_metadata_document, default_tree, envelope_key_b64, make_bundle, ) # --- constants ----------------------------------------------------------------------------------- REPO_ROOT = Path(__file__).resolve().parent.parent FIXTURES_DIR = Path(__file__).resolve().parent / 'fixtures' STATUS_TABLE_DDL = FIXTURES_DIR / 'experiment_run.sql' STATUS_TABLE_DDL_V12 = FIXTURES_DIR / 'experiment_run_v12.sql' GOLDEN_BUNDLE = REPO_ROOT / 'tests' / 'fixtures' / 'bundle' / 'golden_model_v3.sientia' GOLDEN_BUNDLE_PASSWORD = 'Import1Model!Fixture' BFF_DATABASE = 'sientia-core-mlops-bff' BFF_USER = 'bff' BFF_PASSWORD = 'bff' # Lowercase and hyphens only: `chk_experiment_run_bucket_name_pattern` on the real table. IMPORT_BUCKET = 'sientia-import' IMPORT_PREFIX = 'imported_models/' MODELS_COLLECTION = 'models' MONGO_DATABASE = 'sientia' IMPORT_TASK_QUEUE = 'import-model-e2e-queue' # Deliberately gapped, and deliberately strings: the baseline migration stores `id` as a string and # `max(int(id)) + 1` must not be confused with a count or with a lexicographic maximum ("9" > "30"). SEEDED_MODEL_IDS = ('1', '2', '7') MLFLOW_HEALTH_TIMEOUT_SECONDS = 90.0 MLFLOW_HEALTH_POLL_SECONDS = 0.5 IMPORT_ACTIVITY_NAMES = ( 'claim_import_status', 'record_import_status', 'record_import_names', 'record_import_terminal_status', 'report_import_status_write_failure', 'download_import_bundle', 'open_import_bundle', 'create_import_experiment', 'upload_import_artifacts', 'register_import_model_version', 'write_import_model_document', 'cleanup_import_files', 'ensure_import_bundle_retention', ) # --- the BFF Postgres container ------------------------------------------------------------------ @pytest.fixture(scope='session') def import_bff_container() -> Iterator[PostgresContainer]: """The import log database, on a server of its own. Separate from the shared `postgres_container` for two independent reasons, either sufficient: the retry-recovery tests **stop and start this server**, which must not disturb the predictions database every other e2e suite shares session scope with; and "the import touches no Postgres database other than `sientia-core-mlops-bff`" is a stronger statement when the other database lives on a different server. `log_statement=all` is what makes the absence assertions possible: the server log is the only vantage point that sees every client, including the psycopg pool `Storage` holds. """ container = PostgresContainer( 'postgres:15', username=BFF_USER, password=BFF_PASSWORD, dbname=BFF_DATABASE, ).with_command('postgres -c log_statement=all') container.start() try: yield container finally: container.stop() @pytest.fixture(scope='session') def import_bff_engine(import_bff_container: PostgresContainer) -> Iterator[Engine]: """A SQLAlchemy engine on the BFF container. `pool_pre_ping` matters here rather than being habit: two tests stop the container underneath this engine, and a stale pooled connection would fail the *fixture* instead of the assertion. """ engine = create_engine(import_bff_container.get_connection_url(), pool_pre_ping=True) try: yield engine finally: engine.dispose() def _apply_ddl(engine: Engine, files: tuple[Path, ...]) -> None: """Drop `public.experiment_run` and rebuild it from `files`, in order.""" with engine.begin() as connection: connection.execute(text('DROP TABLE IF EXISTS public.experiment_run CASCADE')) for path in files: # One transaction per file, because each file is one migration's worth of statements. with engine.begin() as connection: connection.execute(text(path.read_text(encoding='utf-8'))) @pytest.fixture def import_status_table(import_bff_engine: Engine) -> Engine: """`public.experiment_run` as `V1`…`V12` leave it, rebuilt for this test. Rebuilt rather than truncated: the pre-`V12` variant below has to be able to alternate with this one inside a single session, and `TRUNCATE` cannot take a constraint away. The effect a test sees is the one the spec asks for — the table holds only the row that test inserted. """ _apply_ddl(import_bff_engine, (STATUS_TABLE_DDL, STATUS_TABLE_DDL_V12)) return import_bff_engine @pytest.fixture def import_status_table_pre_v12(import_bff_engine: Engine) -> Engine: """The same table without `V12`: `run_type` present, the import statuses and pairing absent. There is no second copy of the base schema — this applies the same file, and simply does not apply the delta. """ _apply_ddl(import_bff_engine, (STATUS_TABLE_DDL,)) return import_bff_engine InsertStatusRow = Callable[..., int] @pytest.fixture def import_insert_status_row(import_bff_engine: Engine) -> InsertStatusRow: """Insert one import log row the way the BFF will, and return its `id`. Defaults mirror what the frontend sends: `PENDING`, `run_type='IMPORT'`, the object key in `file_name`, and the uploader's digest inside `request_data` — which is where `claim_import_status` looks for it. """ def insert( *, object_key: str, expected_digest: str | None = None, username: str = 'import.e2e@example.com', experiment_name: str = 'import-e2e-placeholder', status: str = 'PENDING', run_type: str = 'IMPORT', bucket_name: str = IMPORT_BUCKET, request_data: dict[str, Any] | None = None, orchestrator_response_data: dict[str, Any] | None = None, ) -> int: payload = request_data if request_data is not None else {} if expected_digest is not None and request_data is None: payload = {'expected_digest': expected_digest} now = datetime.now(UTC).replace(tzinfo=None) statement = text( 'INSERT INTO public.experiment_run ' '(experiment_name, username, status, created_at, updated_at, bucket_name, file_name, ' ' request_data, orchestrator_response_data, run_type) ' 'VALUES (:experiment_name, :username, :status, :created_at, :updated_at, ' ' :bucket_name, :file_name, CAST(:request_data AS JSONB), ' ' CAST(:detail AS JSONB), :run_type) ' 'RETURNING id' ) with import_bff_engine.begin() as connection: return int( connection.execute( statement, { 'experiment_name': experiment_name, 'username': username, 'status': status, 'created_at': now, 'updated_at': now, 'bucket_name': bucket_name, 'file_name': object_key, 'request_data': json.dumps(payload), 'detail': ( json.dumps(orchestrator_response_data) if orchestrator_response_data is not None else None ), 'run_type': run_type, }, ).scalar_one() ) return insert ReadStatusRow = Callable[[int], dict[str, Any] | None] @pytest.fixture def import_read_status_row(import_bff_engine: Engine) -> ReadStatusRow: """Read one import log row as a plain mapping, or `None` when it is gone.""" def read(import_run_id: int) -> dict[str, Any] | None: with import_bff_engine.begin() as connection: row = ( connection.execute( text('SELECT * FROM public.experiment_run WHERE id = :id'), {'id': import_run_id}, ) .mappings() .one_or_none() ) return dict(row) if row is not None else None return read # --- MongoDB ------------------------------------------------------------------------------------ @pytest.fixture(scope='session') def import_mongo_container() -> Iterator[MongoDbContainer]: """A real MongoDB for the `models` collection the import writes its document into.""" container = MongoDbContainer('mongo:7') container.start() try: yield container finally: container.stop() @pytest.fixture(scope='session') def import_mongo_client(import_mongo_container: MongoDbContainer) -> Iterator[MongoClient]: """A client on the Mongo container, for seeding and for reading assertions back.""" client: MongoClient = MongoClient(import_mongo_container.get_connection_url()) try: yield client finally: client.close() @pytest.fixture def import_models_collection(import_mongo_client: MongoClient) -> Any: """The `models` collection, dropped and reseeded with gapped ids and a unique index on `id`. The index is the assertion's teeth: without it "the id continues the collection" and "a name already present is left alone" would both pass against a collection that quietly took a duplicate. """ collection = import_mongo_client[MONGO_DATABASE][MODELS_COLLECTION] collection.drop() collection.insert_many( [ { 'id': model_id, 'name': f'seeded_model_{model_id}', 'active': True, 'model_config': { 'transform_flavor': 'sklearn', 'predict_flavor': 'joblib', 'retention_minutes': 60, 'target': 'SE', }, } for model_id in SEEDED_MODEL_IDS ] ) collection.create_index('id', unique=True) return collection # --- MinIO -------------------------------------------------------------------------------------- @pytest.fixture def import_bucket(minio_container) -> Iterator[Any]: """The import bucket in the shared MinIO container, with its prefix emptied for this test. The bucket is kept rather than recreated — a bucket carries the retention lifecycle rule the import sets, and dropping it between tests would make that rule untestable. """ client = minio_container.get_client() if not client.bucket_exists(IMPORT_BUCKET): client.make_bucket(IMPORT_BUCKET) _empty_prefix(client, IMPORT_BUCKET, IMPORT_PREFIX) yield client _empty_prefix(client, IMPORT_BUCKET, IMPORT_PREFIX) def _empty_prefix(client: Any, bucket: str, prefix: str) -> None: """Remove every object under `prefix`, so a leftover cannot satisfy the next test.""" for obj in client.list_objects(bucket, prefix=prefix, recursive=True): client.remove_object(bucket, obj.object_name) UploadBundle = Callable[..., tuple[str, str]] @pytest.fixture def import_upload_bundle(import_bucket: Any) -> UploadBundle: """Upload bundle bytes under the import prefix; return `(object_key, hex_digest)`. The digest is computed over the bytes handed in, which is what the frontend promises and what gate 1 holds it to. A test that wants the digest to *disagree* passes its own. """ import io def upload(payload: bytes, name: str | None = None) -> tuple[str, str]: object_name = name or f'e2e_{uuid.uuid4().hex}.sientia' object_key = f'{IMPORT_PREFIX}{object_name}' import_bucket.put_object( IMPORT_BUCKET, object_key, io.BytesIO(payload), length=len(payload), content_type='application/octet-stream', ) return object_key, hashlib.sha256(payload).hexdigest() return upload # --- building bundles ---------------------------------------------------------------------------- @pytest.fixture def import_envelope() -> Callable[[str], str]: """Wrap a bundle password the way the frontend will, under the key the worker is configured with. The key is `bundle_factory.DEFAULT_ENVELOPE_KEY`, which `import_activities` puts in `IMPORT_PASSWORD_KEY` — so an envelope built here is one the worker can actually open, and an envelope built under any other key is a decryption-step failure, which is a test of its own. """ def build(password: str = DEFAULT_PASSWORD) -> str: return build_envelope(password, DEFAULT_ENVELOPE_KEY) return build @pytest.fixture def import_bundle_factory(import_names: dict[str, str]) -> Callable[..., bytes]: """Build a `.sientia` bundle carrying *this test's* unique model and experiment names. Every knob the gate tests need is passed straight through to the shared factory — there is no second implementation of the producer's header, KDF or framing here, which is the whole point of reusing `tests/helpers/bundle_factory.py`. The names default to `import_names` so two tests provisioning models in the same session never observe each other's registry state (MLflow is deliberately not reset — design E13). A test that *wants* the collision, like the second-import case, passes the same names twice. """ def build( *, password: str = DEFAULT_PASSWORD, metadata: dict[str, Any] | None = None, parameters: dict[str, Any] | None = None, drop_metadata_field: str | None = None, drop_top_level_key: str | None = None, drop_parameter: str | None = None, files: dict[str, bytes] | None = None, include_metadata: bool = True, symlinks: dict[str, str] | None = None, absolute_entry: str | None = None, traversal_entry: str | None = None, bomb_bytes: int = 0, extra_files: dict[str, bytes] | None = None, **header_kwargs: Any, ) -> bytes: document = build_metadata_document( metadata={ 'model_name': import_names['model_name'], 'experiment_name': import_names['experiment_name'], **(metadata or {}), }, parameters={'model_name': import_names['model_name'], **(parameters or {})}, drop_metadata_field=drop_metadata_field, drop_top_level_key=drop_top_level_key, drop_parameter=drop_parameter, ) spec = ZipSpec( files=default_tree() if files is None else files, document=document, include_metadata=include_metadata, symlinks=symlinks or {}, absolute_entry=absolute_entry, traversal_entry=traversal_entry, bomb_bytes=bomb_bytes, extra_files=extra_files or {}, ) return make_bundle(spec, password, **header_kwargs) return build # --- the MLflow tracking server ----------------------------------------------------------------- def _free_port() -> int: """A loopback port nothing is listening on, taken from an ephemeral bind.""" with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: sock.bind(('127.0.0.1', 0)) return int(sock.getsockname()[1]) @pytest.fixture(scope='session') def import_mlflow_server(tmp_path_factory: pytest.TempPathFactory) -> Iterator[dict[str, Any]]: """`mlflow server` from *this* venv, on a free loopback port. A subprocess rather than a container because the venv pin **is** the production pin (2.10.1, the same version the producer pins), so this tests the exact client/server pair that will run in production — no image tag to drift out of step with `requirements.txt`. sqlite for the backend store because the model registry is not available on a file-backed tracking store, and `create_model_version` is half of what the provisioning activities do. `--serve-artifacts` sends `log_artifacts` over HTTP through the server, as it goes in production, while leaving the uploaded tree readable straight off the local filesystem. Yields the host, the port and the artifacts directory. A server that never answers `/health` fails *here*, with its own stderr attached, rather than making the first test fail obscurely. """ root = tmp_path_factory.mktemp('import_mlflow') artifacts = root / 'mlartifacts' artifacts.mkdir() log_path = root / 'mlflow-server.log' port = _free_port() host = '127.0.0.1' tracking_uri = f'http://{host}:{port}' server: dict[str, Any] = { 'host': f'http://{host}', 'port': str(port), 'tracking_uri': tracking_uri, 'artifacts_dir': artifacts, 'root': root, 'log_path': log_path, 'process': _launch_mlflow(root, host, port, artifacts, log_path), } try: _wait_for_mlflow(server['process'], tracking_uri, log_path) yield server finally: _terminate_group(server['process']) def _launch_mlflow( root: Path, host: str, port: int, artifacts: Path, log_path: Path ) -> subprocess.Popen: """Start the server process. Separate from the fixture so one test can restart it. The sqlite file and the artifacts directory are reused across a restart on purpose: a test that stops the server to fail one activity must not also erase what earlier tests provisioned. """ with log_path.open('ab') as log_handle: return subprocess.Popen( # noqa: S603 - fixed argv, no shell, all values built here [ sys.executable, '-m', 'mlflow', 'server', '--host', host, '--port', str(port), '--backend-store-uri', f'sqlite:///{root / "mlflow.db"}', '--artifacts-destination', str(artifacts), '--serve-artifacts', ], stdout=log_handle, stderr=subprocess.STDOUT, cwd=str(root), # Its own process group: `mlflow server` is gunicorn, so it has children, and killing # only the parent would leave workers holding the port for the rest of the session. start_new_session=True, ) @pytest.fixture def import_mlflow_stopped(import_mlflow_server: dict[str, Any]) -> Any: """A context manager that stops the tracking server and **always** brings it back. This is design E7's second preference in its purest form: no patch, no mock — the dependency is genuinely gone, which is the thing that actually happens in production when a pod is replaced. The restart is in a `finally` because the server is session-scoped: a test that left it down would fail every test after it, and the failure would point at the wrong place. """ @contextlib.contextmanager def stopped() -> Iterator[None]: _terminate_group(import_mlflow_server['process']) try: yield finally: import_mlflow_server['process'] = _launch_mlflow( import_mlflow_server['root'], '127.0.0.1', int(import_mlflow_server['port']), import_mlflow_server['artifacts_dir'], import_mlflow_server['log_path'], ) _wait_for_mlflow( import_mlflow_server['process'], import_mlflow_server['tracking_uri'], import_mlflow_server['log_path'], ) return stopped def _wait_for_mlflow(process: subprocess.Popen, tracking_uri: str, log_path: Path) -> None: """Poll `/health` until it answers, or fail with what the server actually said.""" deadline = time.monotonic() + MLFLOW_HEALTH_TIMEOUT_SECONDS while time.monotonic() < deadline: if process.poll() is not None: _terminate_group(process) raise RuntimeError( f'mlflow server exited with {process.returncode} before answering /health:\n' f'{log_path.read_text(encoding="utf-8", errors="replace")}' ) try: if requests.get(f'{tracking_uri}/health', timeout=2).status_code == 200: return except requests.RequestException: pass time.sleep(MLFLOW_HEALTH_POLL_SECONDS) _terminate_group(process) raise RuntimeError( f'mlflow server did not answer /health within {MLFLOW_HEALTH_TIMEOUT_SECONDS}s:\n' f'{log_path.read_text(encoding="utf-8", errors="replace")}' ) def _terminate_group(process: subprocess.Popen) -> None: """Stop the server and every gunicorn worker it forked.""" if process.poll() is not None: return with contextlib.suppress(ProcessLookupError, PermissionError): os.killpg(os.getpgid(process.pid), signal.SIGTERM) try: process.wait(timeout=15) except subprocess.TimeoutExpired: with contextlib.suppress(ProcessLookupError, PermissionError): os.killpg(os.getpgid(process.pid), signal.SIGKILL) process.wait(timeout=15) # --- per-test naming ---------------------------------------------------------------------------- @pytest.fixture def import_names(request: pytest.FixtureRequest) -> dict[str, str]: """Experiment and model names unique to this test. MLflow state is deliberately not reset between tests (design E13): deleting a registered model leaves a tombstone `get_registered_model` still finds, which would make "second import of the same model" read as a leak rather than as the thing it is testing. Unique names per test buy the same isolation without the tombstones. """ token = uuid.uuid4().hex[:8] stem = request.node.name.replace('[', '_').replace(']', '')[:24] return { 'token': token, 'model_name': f'm_{stem}_{token}'[:40], 'experiment_name': f'e_{stem}_{token}'[:40], } # --- the import Activities ---------------------------------------------------------------------- @pytest.fixture def import_limits_override() -> dict[str, str]: """Extra import environment variables for this test, applied last so they win. Empty by default. A test that needs a ceiling it can actually cross — gate 1's object size is 1 GiB, and uploading a gigabyte to prove a limit is not a test anybody runs twice — overrides this fixture inside its own class, so the narrowed ceiling is scoped to the cases that need it and every other test keeps the production defaults. """ return {} @pytest_asyncio.fixture async def import_activities( monkeypatch: pytest.MonkeyPatch, import_limits_override: dict[str, str], tmp_path: Path, postgres_container, import_bff_container: PostgresContainer, import_mongo_container: MongoDbContainer, import_mlflow_server: dict[str, Any], import_bucket: Any, minio_container, mock_logger, notification_handler, ) -> Any: """An `Activities` whose import path talks to nothing fake. Note what is *not* in the fixture's parameter list: `patch_mlflow`, `patch_minio_repository` and `patch_create_engine`. Composing any of them would replace exactly the thing under test, and because `mlflow.set_tracking_uri` is process-global the substitution would be silent — so the live tracking URI is asserted below rather than assumed. The import configuration is built by the production `build_import_config` / `build_import_status_config` from environment variables set here, so the variable *names* are covered too and not just the dictionaries they produce. """ from laborious.utils.connectors_config import build_import_config, build_import_status_config minio_port = minio_container.get_exposed_port(9000) monkeypatch.setenv('IMPORT_PASSWORD_KEY', envelope_key_b64(DEFAULT_ENVELOPE_KEY)) monkeypatch.setenv('IMPORT_BUNDLE_BUCKET', IMPORT_BUCKET) monkeypatch.setenv('IMPORT_BUNDLE_PREFIX', IMPORT_PREFIX) monkeypatch.setenv('IMPORT_MODELS_COLLECTION', MODELS_COLLECTION) monkeypatch.setenv('IMPORT_WORK_DIR', str(tmp_path / 'import-work')) monkeypatch.setenv('IMPORT_STATUS_DB_HOST', import_bff_container.get_container_host_ip()) monkeypatch.setenv('IMPORT_STATUS_DB_PORT', str(import_bff_container.get_exposed_port(5432))) monkeypatch.setenv('IMPORT_STATUS_DB_NAME', BFF_DATABASE) monkeypatch.setenv('IMPORT_STATUS_DB_USER', BFF_USER) monkeypatch.setenv('IMPORT_STATUS_DB_PASSWORD', BFF_PASSWORD) for name, value in import_limits_override.items(): monkeypatch.setenv(name, value) activities = Activities( # The predictions database, which the import must never touch. It is pointed at a real # server on purpose: "no statement reached it" is only worth asserting when there is a log # that would have shown one. postgres_config={ 'host': 'localhost', 'port': postgres_container.get_exposed_port(5432), 'user': 'test', 'password': 'test', 'dbname': 'test', 'min_connections': 1, 'max_connections': 5, }, mlflow_config={ 'host': import_mlflow_server['host'], 'port': import_mlflow_server['port'], 'username': 'e2e', 'password': 'e2e', }, minio_config={ 'endpoint_url': f'localhost:{minio_port}', 'access_key': 'minioadmin', 'secret_key': 'minioadmin', 'default_bucket': IMPORT_BUCKET, 'retention_hours': 24, 'secure': False, }, opc_config={}, pi_web_api_config={ 'base_url': 'http://localhost:8080', 'auth_type': 'bearer', 'auth_token': 'e2e', }, logger=mock_logger, notification_handler=notification_handler, mongo_config={ 'connection_string': import_mongo_container.get_connection_url(), 'database_name': MONGO_DATABASE, }, import_config=build_import_config(), import_status_config=build_import_status_config(), ) # If a mocked-MLflow fixture ever leaks into an import test, this is where it stops. assert mlflow.get_tracking_uri() == import_mlflow_server['tracking_uri'], ( f'MLflow tracking URI is {mlflow.get_tracking_uri()!r}, not the test server ' f'{import_mlflow_server["tracking_uri"]!r} — a mocked MLflow fixture has leaked in' ) try: yield activities finally: await activities.shutdown() @pytest_asyncio.fixture async def import_worker(temporal_test_env, import_activities: Any) -> Any: """An `ImportModel` worker with the thirteen import activities, on the import queue.""" async with Worker( temporal_test_env.client, task_queue=IMPORT_TASK_QUEUE, workflows=[ImportModel], activities=[getattr(import_activities, name) for name in IMPORT_ACTIVITY_NAMES], ) as worker: yield worker # --- running an import -------------------------------------------------------------------------- @pytest.fixture def import_build_input() -> Callable[..., dict[str, Any]]: """The workflow input, as the frontend assembles it.""" def build( *, import_run_id: int, object_key: str, expected_digest: str | None, password_envelope: str, bucket: str | None = IMPORT_BUCKET, **extra: Any, ) -> dict[str, Any]: payload: dict[str, Any] = { 'import_run_id': import_run_id, 'object_key': object_key, 'password_envelope': password_envelope, } if expected_digest is not None: payload['expected_digest'] = expected_digest if bucket is not None: payload['bucket'] = bucket payload.update(extra) return payload return build @pytest.fixture def import_execute(temporal_test_env, import_worker: Any) -> Callable[..., Any]: """Run one import to completion and hand back the handle, whatever the outcome. The handle rather than the result, because half this suite asserts on *how* a run failed — the error type Temporal recorded, the event history — and `execute_workflow` would have thrown that away. """ async def execute(payload: dict[str, Any], workflow_id: str | None = None) -> WorkflowHandle: handle = await temporal_test_env.client.start_workflow( ImportModel.run, payload, id=workflow_id or f'import-e2e-{uuid.uuid4().hex}', task_queue=IMPORT_TASK_QUEUE, ) with contextlib.suppress(Exception): # The exception is the assertion's subject in the negative cases; the caller reads it # back off the handle so a failure here is never mistaken for a harness error. await handle.result() return handle return execute @pytest.fixture def import_history_events() -> Callable[[WorkflowHandle], Any]: """Every event of a run's history, flattened to the text the secrecy assertions scan. Returns `(events, blob)`: the event list for structural assertions, and one string holding every payload in the history for "this value appears nowhere" assertions. """ async def fetch(handle: WorkflowHandle) -> tuple[list[Any], str]: events = [event async for event in handle.fetch_history_events()] return events, '\n'.join(str(event) for event in events) return fetch # --- reading a failure --------------------------------------------------------------------------- # # Temporal wraps what actually went wrong: `WorkflowFailureError` -> `ActivityError` -> # `ApplicationError`, and only the innermost one carries the sentence and the declared error type. # `str(exception)` on the outermost is "Workflow execution failed", which is true and useless, so # every negative assertion in this suite goes through these two functions instead. def failure_chain(error: BaseException) -> list[BaseException]: """The exception and everything it was caused by, outermost first.""" chain: list[BaseException] = [] current: BaseException | None = error while current is not None and not any(current is seen for seen in chain): chain.append(current) current = current.__cause__ return chain def failure_text(error: BaseException) -> str: """Every link of the chain as `Type: message`, joined — what a substring assertion reads.""" return '\n'.join(f'{type(part).__name__}: {part}' for part in failure_chain(error)) def failure_types(error: BaseException) -> set[str]: """Every Python class name in the chain, plus every `ApplicationError.type` it declares. The declared `type` is the one that survives serialisation into the event history, so it is what "the run failed with the not-claimable error type" actually means. """ names: set[str] = set() for part in failure_chain(error): names.add(type(part).__name__) declared = getattr(part, 'type', None) if isinstance(declared, str): names.add(declared) return names # --- the absence assertion ---------------------------------------------------------------------- @pytest.fixture def import_assert_nothing_provisioned( import_mlflow_server: dict[str, Any], import_models_collection: Any, import_bucket: Any, tmp_path: Path, ) -> Callable[..., None]: """Assert a rejected import left nothing behind, in one call over all four places. One helper rather than four assertions per test, because "nothing was provisioned" is a single claim and a test that checks three of its four halves reads as if it checked all four. """ def assert_nothing( *, experiment_name: str, model_name: str, object_key: str | None = None, seeded_ids: tuple[str, ...] = SEEDED_MODEL_IDS, ) -> None: client = mlflow.tracking.MlflowClient(tracking_uri=import_mlflow_server['tracking_uri']) experiment = client.get_experiment_by_name(experiment_name) assert experiment is None, f'experiment {experiment_name!r} was created' versions = client.search_model_versions(f"name='{model_name}'") assert not list(versions), f'registry holds versions of {model_name!r}' documents = list(import_models_collection.find({}, {'_id': False})) assert {str(document['id']) for document in documents} == set(seeded_ids), ( f'the models collection changed: {documents}' ) assert not [d for d in documents if d.get('name') == model_name], ( f'a model document for {model_name!r} was written' ) # The uploaded object itself stays — the frontend put it there and retention removes it. # What must not be there is anything the *import* wrote under the prefix. stored = { obj.object_name for obj in import_bucket.list_objects( IMPORT_BUCKET, prefix=IMPORT_PREFIX, recursive=True ) } assert stored <= ({object_key} if object_key else set()), ( f'the import wrote objects under {IMPORT_PREFIX}: {stored}' ) work_dir = tmp_path / 'import-work' leftovers = list(work_dir.rglob('*')) if work_dir.exists() else [] assert not leftovers, f'the worker filesystem still holds {leftovers}' return assert_nothing # --- reading the golden fixture ----------------------------------------------------------------- @pytest.fixture(scope='session') def import_golden_bundle() -> dict[str, Any]: """The producer-emitted bundle: its bytes, its digest and the password it was written with. This is the only place in the suite that pays the producer's real Argon2id cost (256 MiB, 3 passes). Everything else runs the factory's cheap cost, which is the right split — the reader honours whatever the header declares, up to its ceiling. """ payload = GOLDEN_BUNDLE.read_bytes() return { 'bytes': payload, 'digest': hashlib.sha256(payload).hexdigest(), 'password': GOLDEN_BUNDLE_PASSWORD, 'path': GOLDEN_BUNDLE, }