vitor-aignosi
192834bde3
QTZPOC-20: Refactor model import workflow, introduce atomic model ID generation, and clarify MinIO expiry. Fuses bundle download, decryption, MLflow experiment creation, and artifact upload into a single activity to prevent cross-worker issues. Implements an atomic counter for MongoDB model IDs, replacing the max(id) + 1 scan to avoid collisions during concurrent imports. Adds comprehensive documentation for MinIO object expiry mechanisms and refines lifecycle rule application. Improves error handling with clearer responsibility classification and centralizes MLflow repository injection.
1344 lines
54 KiB
Python
1344 lines
54 KiB
Python
"""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 asyncio
|
||
import contextlib
|
||
import hashlib
|
||
import json
|
||
import os
|
||
import re
|
||
import signal
|
||
import socket
|
||
import subprocess
|
||
import sys
|
||
import time
|
||
import uuid
|
||
from collections.abc import Callable, Iterator
|
||
from datetime import UTC, datetime, timedelta
|
||
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.common import RetryPolicy
|
||
from temporalio.worker import Worker
|
||
from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner, SandboxRestrictions
|
||
from testcontainers.community.mongodb import MongoDbContainer
|
||
from testcontainers.community.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'
|
||
# Must track `IMPORT_MODEL_ID_COUNTERS_COLLECTION`'s default in connectors_config.py — the harness
|
||
# never sets that env var, so production's default is what the worker under test actually reads from.
|
||
MODEL_ID_COUNTERS_COLLECTION = 'model_id_counters'
|
||
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
|
||
MONGO_SERVER_SELECTION_MS = 5000
|
||
|
||
IMPORT_ACTIVITY_NAMES = (
|
||
'claim_import_status',
|
||
'record_import_status',
|
||
'record_import_names',
|
||
'record_import_terminal_status',
|
||
'report_import_status_write_failure',
|
||
'report_import_pipeline_failure',
|
||
'prepare_and_upload_import_bundle',
|
||
'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.
|
||
|
||
The host port is **bound explicitly** rather than left ephemeral. Docker re-allocates an
|
||
ephemeral published port when a container is started again, so a stop/start would move the
|
||
database to a new address and the retry-recovery tests would be asserting that a *different*
|
||
server came up. A fixed binding is what makes "it came back" mean what it says.
|
||
"""
|
||
container = (
|
||
PostgresContainer(
|
||
'postgres:15',
|
||
username=BFF_USER,
|
||
password=BFF_PASSWORD,
|
||
dbname=BFF_DATABASE,
|
||
)
|
||
.with_command('postgres -c log_statement=all')
|
||
.with_bind_ports(5432, _free_port())
|
||
)
|
||
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.
|
||
|
||
New ids no longer come from scanning this collection — `_next_model_document_id` reads an
|
||
atomic counter in a separate `model_id_counters` collection instead (a different database, no
|
||
less), so it survives this drop untouched. Reset it here too, to the same seeded high-water
|
||
mark, or the next id handed out silently carries over from whatever an earlier test in the same
|
||
run already claimed.
|
||
"""
|
||
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)
|
||
|
||
counters = import_mongo_client[MONGO_DATABASE][MODEL_ID_COUNTERS_COLLECTION]
|
||
counters.replace_one(
|
||
{'_id': MODELS_COLLECTION},
|
||
{'_id': MODELS_COLLECTION, 'seq': max(int(seeded) for seeded in SEEDED_MODEL_IDS)},
|
||
upsert=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={
|
||
# A short server-selection timeout so the "MongoDB is down" case fails in seconds
|
||
# instead of spending pymongo's 30-second default. Everything else about the client is
|
||
# the production one.
|
||
'connection_string': (
|
||
f'{import_mongo_container.get_connection_url()}'
|
||
f'?serverSelectionTimeoutMS={MONGO_SERVER_SELECTION_MS}'
|
||
),
|
||
'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.
|
||
|
||
`bson` and `pymongo` are declared **passthrough** on the sandbox. Without that, the workflow
|
||
module's `imports_passed_through()` block pulls in `Activities` → `MongoDBRepository` → `pymongo`
|
||
during the first workflow activation, and `pymongo` loads its BSON codecs lazily
|
||
(`bson.objectid`, `bson.timestamp`, `bson.decimal128`, …). Each one is an import the sandbox has
|
||
not seen, so it warns "Module bson.X was imported after initial workflow load" — about twenty
|
||
lines per test that touches MongoDB.
|
||
|
||
Declaring them passthrough is the supported fix and the accurate one: these modules are pure
|
||
serialisation code with no workflow-relevant global state, which is exactly what passthrough is
|
||
for. Eagerly importing them on the host does **not** help — the sandbox judges imports against
|
||
its own module table during an activation, not against the host's `sys.modules`.
|
||
"""
|
||
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],
|
||
workflow_runner=SandboxedWorkflowRunner(
|
||
restrictions=SandboxRestrictions.default.with_passthrough_modules('bson', 'pymongo')
|
||
),
|
||
) as worker:
|
||
yield worker
|
||
|
||
|
||
# --- the real-clock environment, for the retry-envelope tests ------------------------------------
|
||
#
|
||
# Design E9: everything else in this suite runs on the shared time-skipping environment. Two claims
|
||
# cannot: "the database came back *during* the retry envelope" needs the restart to occupy real
|
||
# time, and "the failure landed on the *first* attempt" is an assertion about elapsed attempts. Both
|
||
# run here instead, and this is the one place the harness deliberately weakens fidelity — so it says
|
||
# so, rather than hiding it in a fixture.
|
||
|
||
|
||
@pytest.fixture
|
||
def import_status_retry_shortening() -> dict[str, Any]:
|
||
"""How much to shorten the status retry policies for a real-clock test.
|
||
|
||
The **real** values (10 attempts, 1 s doubling to 30 s, a 300 s envelope) stay asserted where
|
||
they belong — unit test 12.1 of `import-model-workflow` pins them. What the tests here prove is
|
||
the recovery *behaviour*, which is identical at any interval, so waiting 151 seconds to observe
|
||
it would be waiting for arithmetic.
|
||
|
||
A constant half-second interval rather than a shortened doubling: it makes the envelope a number
|
||
a reader can multiply out (`attempts × interval`) instead of a geometric series.
|
||
"""
|
||
return {
|
||
'initial_interval': 0.5,
|
||
'backoff_coefficient': 1.0,
|
||
'maximum_attempts': 24, # 12 s of retrying
|
||
'hint_maximum_attempts': 2,
|
||
'envelope_seconds': 25,
|
||
'write_timeout_seconds': 5,
|
||
}
|
||
|
||
|
||
@pytest_asyncio.fixture
|
||
async def import_real_clock_env() -> Any:
|
||
"""A `WorkflowEnvironment` on the real clock."""
|
||
from temporalio.testing import WorkflowEnvironment
|
||
|
||
env = await WorkflowEnvironment.start_local()
|
||
async with env:
|
||
yield env
|
||
|
||
|
||
@pytest_asyncio.fixture
|
||
async def import_worker_real_clock(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
import_real_clock_env: Any,
|
||
import_status_retry_shortening: dict[str, Any],
|
||
import_activities: Any,
|
||
) -> Any:
|
||
"""An `ImportModel` worker on the real clock, with the status retry policies shortened.
|
||
|
||
Two harness choices worth knowing about:
|
||
|
||
- The policies are patched **on the workflow module** before the worker starts, because the
|
||
workflow reads them as module globals when it schedules an activity.
|
||
- The worker runs `UnsandboxedWorkflowRunner`. The sandbox imports the workflow module into a
|
||
namespace of its own, so a patch applied to the host module would simply not be seen and the
|
||
test would silently run the production 151-second envelope. Unsandboxed is what makes the
|
||
patch effective; nothing else about the workflow changes.
|
||
"""
|
||
from temporalio.worker import UnsandboxedWorkflowRunner
|
||
|
||
from laborious.workflows import import_model as workflow_module
|
||
|
||
shortening = import_status_retry_shortening
|
||
monkeypatch.setattr(
|
||
workflow_module,
|
||
'IMPORT_STATUS_RETRY_POLICY',
|
||
RetryPolicy(
|
||
maximum_attempts=int(shortening['maximum_attempts']),
|
||
initial_interval=timedelta(seconds=shortening['initial_interval']),
|
||
backoff_coefficient=shortening['backoff_coefficient'],
|
||
maximum_interval=timedelta(seconds=shortening['initial_interval']),
|
||
non_retryable_error_types=[
|
||
'ImportLogRowNotClaimableError',
|
||
'ImportStatusRejectedError',
|
||
],
|
||
),
|
||
)
|
||
monkeypatch.setattr(
|
||
workflow_module,
|
||
'IMPORT_STATUS_HINT_RETRY_POLICY',
|
||
RetryPolicy(
|
||
maximum_attempts=int(shortening['hint_maximum_attempts']),
|
||
initial_interval=timedelta(seconds=shortening['initial_interval']),
|
||
backoff_coefficient=1.0,
|
||
maximum_interval=timedelta(seconds=shortening['initial_interval']),
|
||
non_retryable_error_types=[
|
||
'ImportLogRowNotClaimableError',
|
||
'ImportStatusRejectedError',
|
||
],
|
||
),
|
||
)
|
||
monkeypatch.setattr(
|
||
workflow_module,
|
||
'STATUS_WRITE_ENVELOPE',
|
||
timedelta(seconds=shortening['envelope_seconds']),
|
||
)
|
||
monkeypatch.setattr(
|
||
workflow_module,
|
||
'STATUS_WRITE_TIMEOUT',
|
||
timedelta(seconds=shortening['write_timeout_seconds']),
|
||
)
|
||
|
||
async with Worker(
|
||
import_real_clock_env.client,
|
||
task_queue=IMPORT_TASK_QUEUE,
|
||
workflows=[ImportModel],
|
||
activities=[getattr(import_activities, name) for name in IMPORT_ACTIVITY_NAMES],
|
||
workflow_runner=UnsandboxedWorkflowRunner(),
|
||
) as worker:
|
||
yield worker
|
||
|
||
|
||
@pytest.fixture
|
||
def import_execute_real_clock(import_real_clock_env: Any, import_worker_real_clock: Any) -> Any:
|
||
"""Start an import on the real-clock environment and hand back the handle."""
|
||
|
||
async def execute(payload: dict[str, Any], workflow_id: str | None = None) -> WorkflowHandle:
|
||
return await import_real_clock_env.client.start_workflow(
|
||
ImportModel.run,
|
||
payload,
|
||
id=workflow_id or f'import-e2e-rc-{uuid.uuid4().hex}',
|
||
task_queue=IMPORT_TASK_QUEUE,
|
||
)
|
||
|
||
return execute
|
||
|
||
|
||
@pytest.fixture
|
||
def import_bff_stopped(import_bff_container: PostgresContainer, import_bff_engine: Engine) -> Any:
|
||
"""An **async** context manager that stops the status database and always brings it back.
|
||
|
||
`docker stop` / `docker start` on the wrapped container, not the testcontainers `stop()`, which
|
||
removes the container. The published port is bound explicitly by `import_bff_container`, so the
|
||
database comes back at the address the worker is already configured with.
|
||
|
||
Async, and every blocking call goes through `asyncio.to_thread`, for a reason that cost an
|
||
afternoon to find: the Temporal worker runs on this test's event loop. A synchronous
|
||
`docker stop` — or a plain `time.sleep` in the test body — starves the worker, so the workflow
|
||
makes no progress at all while the database is down and the test ends up asserting that a
|
||
*stalled* workflow behaved correctly. Nothing retries on a blocked loop.
|
||
|
||
The restart is in a `finally` and the engine is disposed on the way in: pooled connections do not
|
||
survive the server going away, and a stale one would fail the next test rather than this one.
|
||
"""
|
||
|
||
@contextlib.asynccontextmanager
|
||
async def stopped() -> Any:
|
||
wrapped = import_bff_container.get_wrapped_container()
|
||
await asyncio.to_thread(import_bff_engine.dispose)
|
||
await asyncio.to_thread(wrapped.stop, timeout=5)
|
||
try:
|
||
yield
|
||
finally:
|
||
await asyncio.to_thread(wrapped.start)
|
||
await asyncio.to_thread(_wait_for_postgres, import_bff_engine)
|
||
|
||
return stopped
|
||
|
||
|
||
def _wait_for_postgres(engine: Engine) -> None:
|
||
"""Block until the restarted server accepts a connection again."""
|
||
deadline = time.monotonic() + 60
|
||
last_error: Exception | None = None
|
||
while time.monotonic() < deadline:
|
||
try:
|
||
with engine.connect() as connection:
|
||
connection.execute(text('SELECT 1'))
|
||
return
|
||
except Exception as error: # noqa: BLE001 - any failure here means "not up yet"
|
||
last_error = error
|
||
engine.dispose()
|
||
time.sleep(0.5)
|
||
raise RuntimeError(f'the status database did not come back: {last_error!r}')
|
||
|
||
|
||
# --- 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
|
||
|
||
|
||
# --- the statement-log window --------------------------------------------------------------------
|
||
#
|
||
# "The import writes to one table in one database" cannot be asserted from the application side.
|
||
# A SQLAlchemy `before_cursor_execute` listener is blind to the psycopg pool `Storage` holds, which
|
||
# is exactly the client that would betray the claim; `pg_stat_statements` is aggregated, so it
|
||
# cannot be scoped to one import. The server's own log is the only vantage point that sees every
|
||
# client, which is why both containers run with `log_statement=all`.
|
||
|
||
_STATEMENT_MARKERS = ('statement:', 'execute <unnamed>:', 'execute ')
|
||
|
||
# A new log entry starts with the server's timestamp; anything else continues the previous one.
|
||
# Statements are logged across as many lines as they were written on, so without this a multi-line
|
||
# `UPDATE ... WHERE id = ...` would be judged on its first line alone.
|
||
_LOG_ENTRY_START = re.compile(r'^\d{4}-\d{2}-\d{2} ')
|
||
|
||
# Statements that name no *user* relation and are therefore not evidence either way.
|
||
_RELATIONLESS_PREFIXES = (
|
||
'begin',
|
||
'commit',
|
||
'rollback',
|
||
'set ',
|
||
'show ',
|
||
'discard',
|
||
'select 1',
|
||
'select current_database',
|
||
'select current_schema',
|
||
'select pg_backend_pid',
|
||
'select version()',
|
||
)
|
||
|
||
# System catalogs. Every driver reads these on connection setup — psycopg2 asks `pg_type` for its
|
||
# OID map before the application issues anything — so a statement whose only relations are catalogs
|
||
# is connection bookkeeping, not a table the feature chose to touch.
|
||
_CATALOG_IDENTIFIERS = (
|
||
'pg_catalog',
|
||
'pg_type',
|
||
'pg_namespace',
|
||
'pg_range',
|
||
'pg_class',
|
||
'pg_attribute',
|
||
'pg_proc',
|
||
'pg_settings',
|
||
'pg_database',
|
||
'information_schema',
|
||
'typarray',
|
||
)
|
||
|
||
|
||
def _container_log_text(container: Any) -> str:
|
||
"""Everything the container has written to stdout and stderr so far, as one string."""
|
||
streams = container.get_logs()
|
||
parts = streams if isinstance(streams, tuple) else (streams,)
|
||
return '\n'.join(
|
||
part.decode('utf-8', errors='replace') if isinstance(part, bytes) else str(part)
|
||
for part in parts
|
||
if part
|
||
)
|
||
|
||
|
||
def _statements_between(log_text: str, begin: str, end: str) -> list[str]:
|
||
"""Every logged statement between the two marker strings, in order.
|
||
|
||
Anchored on markers rather than on timestamps: the log is shared with every other client, and a
|
||
timestamp window would either include their statements or race with the clock.
|
||
"""
|
||
try:
|
||
start_index = log_text.index(begin)
|
||
except ValueError as error:
|
||
raise AssertionError(f'the opening marker {begin!r} is not in the log') from error
|
||
try:
|
||
end_index = log_text.index(end, start_index)
|
||
except ValueError as error:
|
||
raise AssertionError(f'the closing marker {end!r} is not in the log') from error
|
||
|
||
# Expand to whole lines. Slicing on the marker offsets themselves would cut the marker's own log
|
||
# line in half at both ends, and the tail half of the closing one reads like a real statement.
|
||
region_start = log_text.find('\n', start_index)
|
||
region_start = 0 if region_start == -1 else region_start + 1
|
||
region_end = log_text.rfind('\n', start_index, end_index)
|
||
region = log_text[region_start:region_end] if region_end > region_start else ''
|
||
|
||
statements: list[str] = []
|
||
collecting = False
|
||
for line in region.splitlines():
|
||
if _LOG_ENTRY_START.match(line):
|
||
collecting = False
|
||
for marker in _STATEMENT_MARKERS:
|
||
position = line.find(marker)
|
||
if position != -1:
|
||
statements.append(line[position + len(marker) :].strip())
|
||
collecting = True
|
||
break
|
||
elif collecting and statements:
|
||
# A continuation line of the statement above it.
|
||
statements[-1] = f'{statements[-1]} {line.strip()}'
|
||
return [statement for statement in statements if statement and begin not in statement]
|
||
|
||
|
||
_SINGLE_QUOTED = re.compile(r"'(?:[^']|'')*'")
|
||
|
||
|
||
def strip_literals(statement: str) -> str:
|
||
"""The statement with every single-quoted literal removed.
|
||
|
||
Necessary for any "this relation name does not appear" assertion: the import writes the bundle's
|
||
own `model_project` into the JSON detail, so the *value* `factory_project` would otherwise read
|
||
as a reference to a `project` table. A relation name appears outside a literal or not at all.
|
||
"""
|
||
return _SINGLE_QUOTED.sub("''", statement)
|
||
|
||
|
||
def names_no_relation(statement: str) -> bool:
|
||
"""True when the statement touches no *user* relation.
|
||
|
||
Two kinds: transaction and session bookkeeping (`BEGIN`, `SET`, `SELECT 1`), and the system
|
||
catalogs every driver reads when it connects. Neither is a table the feature chose to write.
|
||
"""
|
||
lowered = statement.strip().lower()
|
||
if any(lowered.startswith(prefix) for prefix in _RELATIONLESS_PREFIXES):
|
||
return True
|
||
return any(identifier in lowered for identifier in _CATALOG_IDENTIFIERS)
|
||
|
||
|
||
@pytest.fixture
|
||
def import_log_window(
|
||
import_bff_engine: Engine,
|
||
import_bff_container: PostgresContainer,
|
||
postgres_engine: Any,
|
||
postgres_container: Any,
|
||
) -> Any:
|
||
"""A context manager bracketing an import with markers on **both** Postgres servers.
|
||
|
||
On exit it fills the dict it yielded with `bff` and `predictions` — the statements each server
|
||
logged between the markers. A marker is written to both servers so each window is bounded on the
|
||
server it is about, rather than one server's window being inferred from the other's clock.
|
||
"""
|
||
|
||
@contextlib.contextmanager
|
||
def window() -> Iterator[dict[str, list[str]]]:
|
||
token = uuid.uuid4().hex[:12]
|
||
begin = f'e2e_window_begin_{token}'
|
||
end = f'e2e_window_end_{token}'
|
||
captured: dict[str, list[str]] = {}
|
||
|
||
_emit_marker(import_bff_engine, begin)
|
||
_emit_marker(postgres_engine, begin)
|
||
try:
|
||
yield captured
|
||
finally:
|
||
_emit_marker(import_bff_engine, end)
|
||
_emit_marker(postgres_engine, end)
|
||
captured['bff'] = _statements_between(
|
||
_container_log_text(import_bff_container), begin, end
|
||
)
|
||
captured['predictions'] = _statements_between(
|
||
_container_log_text(postgres_container), begin, end
|
||
)
|
||
|
||
return window
|
||
|
||
|
||
def _emit_marker(engine: Engine, marker: str) -> None:
|
||
"""A statement whose only purpose is to be findable in the server's log."""
|
||
with engine.begin() as connection:
|
||
connection.execute(text(f"SELECT '{marker}' AS marker")) # noqa: S608 - a generated uuid
|
||
|
||
|
||
# --- 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,
|
||
}
|