QTZPOC-20
QTZPOC-20: Implement .sientia model bundle import workflow Introduces a new Temporal workflow and activities to import .sientia model bundles. This enables: - Secure decryption and validation of bundles through a seven-gate process. - Provisioning imported models in MLflow and updating the MongoDB model listing. - Comprehensive status tracking in the public.experiment_run table with new IMPORT_SUCCESS/IMPORT_ERROR states. - Dedicated metrics, robust error handling, and configurable settings. - Includes updated e2e database fixtures and extensive unit/integration tests for bundle format, reader, and overall workflow.
This commit is contained in:
0
tests/laborious/utils/bundle/__init__.py
Normal file
0
tests/laborious/utils/bundle/__init__.py
Normal file
34
tests/laborious/utils/bundle/conftest.py
Normal file
34
tests/laborious/utils/bundle/conftest.py
Normal file
@@ -0,0 +1,34 @@
|
||||
"""Shared fixtures for the `.sientia` bundle tests."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
FIXTURE_DIR = Path(__file__).resolve().parents[3] / 'fixtures' / 'bundle'
|
||||
GOLDEN_BUNDLE = FIXTURE_DIR / 'golden_model_v3.sientia'
|
||||
GOLDEN_METADATA = FIXTURE_DIR / 'golden_model_v3.metadata.json'
|
||||
# Recorded in tests/fixtures/bundle/README.md; the fixture protects nothing.
|
||||
GOLDEN_PASSWORD = 'Import1Model!Fixture'
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def golden_bundle() -> Path:
|
||||
"""The `.sientia` file the real producer emitted."""
|
||||
assert GOLDEN_BUNDLE.is_file(), (
|
||||
f'golden fixture missing: {GOLDEN_BUNDLE} — regenerate it with '
|
||||
'tests/fixtures/bundle/regenerate.py'
|
||||
)
|
||||
return GOLDEN_BUNDLE
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def golden_password() -> str:
|
||||
return GOLDEN_PASSWORD
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def work_dir(tmp_path: Path) -> Path:
|
||||
"""A per-test stand-in for `IMPORT_WORK_DIR/<workflow_id>`."""
|
||||
root = tmp_path / 'workflow-id'
|
||||
root.mkdir()
|
||||
return root
|
||||
227
tests/laborious/utils/bundle/test_format.py
Normal file
227
tests/laborious/utils/bundle/test_format.py
Normal file
@@ -0,0 +1,227 @@
|
||||
"""The wire format: header parsing, ceilings, digest coverage and the fixture contract."""
|
||||
|
||||
import ast
|
||||
import hashlib
|
||||
from pathlib import Path
|
||||
|
||||
import nacl.bindings as sodium
|
||||
import pytest
|
||||
|
||||
from laborious.utils.bundle import format as fmt
|
||||
from tests.helpers import bundle_factory as factory
|
||||
|
||||
|
||||
def _write(tmp_path: Path, payload: bytes, name: str = 'bundle.sientia') -> Path:
|
||||
target = tmp_path / name
|
||||
target.write_bytes(payload)
|
||||
return target
|
||||
|
||||
|
||||
# --------------------------------------------------------------- constants
|
||||
|
||||
|
||||
def test_header_size_is_140():
|
||||
assert fmt.HEADER_SIZE == 140
|
||||
assert fmt.HEADER_STRUCT.format == '>8sBBB16sQQB32s64s'
|
||||
|
||||
|
||||
def test_libsodium_literals_match_the_binding():
|
||||
"""The literals in `format.py` keep the module import-free; these assertions keep them true."""
|
||||
assert fmt.STREAM_HEADER_SIZE == sodium.crypto_secretstream_xchacha20poly1305_HEADERBYTES
|
||||
assert fmt.ABYTES == sodium.crypto_secretstream_xchacha20poly1305_ABYTES
|
||||
assert fmt.KEY_SIZE == sodium.crypto_secretstream_xchacha20poly1305_KEYBYTES
|
||||
|
||||
|
||||
def test_producer_constants_match():
|
||||
assert fmt.MAGIC == b'SIENTIA1'
|
||||
assert fmt.CHUNK_SIZE == 65536
|
||||
assert fmt.MAX_KDF_MEMLIMIT == 268435456
|
||||
assert fmt.MAX_KDF_OPSLIMIT == 3
|
||||
assert fmt.SALT_SIZE == 16
|
||||
assert fmt.DIGEST_SIZE == 32
|
||||
assert fmt.RESERVED_SIGNATURE_SIZE == 64
|
||||
|
||||
|
||||
def test_prefix_size_is_65717():
|
||||
"""The pre-upload prefix is computed from constants, never read from a file."""
|
||||
assert fmt.PREFIX_SIZE == 140 + 24 + (65536 + 17)
|
||||
assert fmt.PREFIX_SIZE == 65717
|
||||
|
||||
|
||||
def test_digest_offset_skips_both_headers():
|
||||
assert fmt.DIGEST_OFFSET == 164
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- parse
|
||||
|
||||
|
||||
def test_parses_a_producer_emitted_header(golden_bundle):
|
||||
header = fmt.read_header(golden_bundle)
|
||||
|
||||
assert header.format_version == 1
|
||||
assert header.aead_id == fmt.AEAD_ID_XCHACHA20POLY1305_SECRETSTREAM
|
||||
assert header.kdf_id == fmt.KDF_ID_ARGON2ID
|
||||
assert len(header.salt) == 16
|
||||
assert header.kdf_memlimit_bytes == 268435456
|
||||
assert header.kdf_opslimit == 3
|
||||
assert header.kdf_parallelism == 1
|
||||
assert len(header.ciphertext_digest) == 32
|
||||
assert header.reserved_signature == bytes(64)
|
||||
|
||||
|
||||
def test_wrong_magic_rejected(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(magic=b'NOTSIENT'))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='expected bundle marker'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_unknown_format_version_rejected(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(format_version=2))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='format version'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_unknown_aead_id_rejected(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(aead_id=7))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='encryption identifier'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_unknown_kdf_id_rejected(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(kdf_id=9))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='key derivation identifier'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_memlimit_above_the_ceiling_rejected(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(declared_memlimit=512 * 1024 * 1024))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='memory'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_opslimit_above_the_ceiling_rejected(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(declared_opslimit=4))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='passes'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_zero_cost_parameters_rejected(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(declared_opslimit=0))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='positive'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_reserved_parallelism_must_be_one(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(parallelism=2))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='parallelism'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_non_zero_reserved_signature_rejected(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(reserved_signature=b'\x01' + bytes(63)))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='reserved signature'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_file_shorter_than_header_plus_stream_header_rejected(tmp_path):
|
||||
path = _write(tmp_path, b'SIENTIA1' + bytes(100))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='smaller than'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_parse_header_needs_the_whole_header():
|
||||
with pytest.raises(fmt.BundleFormatError, match='shorter than'):
|
||||
fmt.parse_header(b'SIENTIA1')
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ digest
|
||||
|
||||
|
||||
def test_digest_covers_the_chunk_stream_not_the_stream_header(golden_bundle):
|
||||
"""The producer hashes the chunks before prepending either header — offset 164, not 140."""
|
||||
header = fmt.read_header(golden_bundle)
|
||||
payload = golden_bundle.read_bytes()
|
||||
|
||||
assert hashlib.sha256(payload[164:]).digest() == header.ciphertext_digest
|
||||
assert hashlib.sha256(payload[140:]).digest() != header.ciphertext_digest
|
||||
assert fmt.compute_ciphertext_digest(golden_bundle) == header.ciphertext_digest
|
||||
|
||||
|
||||
def test_verify_ciphertext_digest_accepts_the_fixture(golden_bundle):
|
||||
fmt.verify_ciphertext_digest(golden_bundle, fmt.read_header(golden_bundle))
|
||||
|
||||
|
||||
def test_digest_mismatch_rejected(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(digest_override=bytes(32)))
|
||||
header = fmt.read_header(path)
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='does not match the digest'):
|
||||
fmt.verify_ciphertext_digest(path, header)
|
||||
|
||||
|
||||
def test_truncated_file_fails_the_digest(tmp_path):
|
||||
path = _write(tmp_path, factory.truncate_bundle(factory.make_bundle()))
|
||||
header = fmt.read_header(path)
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='does not match the digest'):
|
||||
fmt.verify_ciphertext_digest(path, header)
|
||||
|
||||
|
||||
def test_file_digest_is_the_whole_object(tmp_path):
|
||||
payload = factory.make_bundle()
|
||||
path = _write(tmp_path, payload)
|
||||
|
||||
assert fmt.file_digest(path) == hashlib.sha256(payload).hexdigest()
|
||||
|
||||
|
||||
# ------------------------------------------------------- factory conformance
|
||||
|
||||
|
||||
def test_factory_header_matches_fixture_layout(golden_bundle, tmp_path):
|
||||
"""The test-only writer must be structurally identical to the producer's output."""
|
||||
produced = fmt.read_header(_write(tmp_path, factory.make_bundle()))
|
||||
golden = fmt.read_header(golden_bundle)
|
||||
|
||||
assert (produced.format_version, produced.aead_id, produced.kdf_id) == (
|
||||
golden.format_version,
|
||||
golden.aead_id,
|
||||
golden.kdf_id,
|
||||
)
|
||||
assert produced.kdf_parallelism == golden.kdf_parallelism
|
||||
assert produced.reserved_signature == golden.reserved_signature
|
||||
assert len(produced.salt) == len(golden.salt)
|
||||
assert len(produced.ciphertext_digest) == len(golden.ciphertext_digest)
|
||||
assert factory.HEADER_STRUCT.format == fmt.HEADER_STRUCT.format
|
||||
assert factory.PRODUCER_KDF_MEMLIMIT_BYTES == golden.kdf_memlimit_bytes
|
||||
assert factory.PRODUCER_KDF_OPSLIMIT == golden.kdf_opslimit
|
||||
|
||||
|
||||
# --------------------------------------------------------------- liftability
|
||||
|
||||
FORBIDDEN_IMPORT_ROOTS = frozenset(
|
||||
{'laborious', 'temporalio', 'mlflow', 'nacl', 'cryptography', 'sqlalchemy', 'pymongo'}
|
||||
)
|
||||
|
||||
|
||||
def test_format_module_imports_nothing_platform_specific():
|
||||
"""`format.py` stays liftable into a shared package: stdlib only, no platform imports."""
|
||||
tree = ast.parse(Path(fmt.__file__).read_text(encoding='utf-8'))
|
||||
roots = set()
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.Import):
|
||||
roots.update(alias.name.split('.')[0] for alias in node.names)
|
||||
elif isinstance(node, ast.ImportFrom) and node.module:
|
||||
roots.add(node.module.split('.')[0])
|
||||
|
||||
assert roots.isdisjoint(FORBIDDEN_IMPORT_ROOTS), roots & FORBIDDEN_IMPORT_ROOTS
|
||||
assert roots == {'__future__', 'hashlib', 'struct', 'dataclasses', 'pathlib', 'typing'}
|
||||
460
tests/laborious/utils/bundle/test_reader.py
Normal file
460
tests/laborious/utils/bundle/test_reader.py
Normal file
@@ -0,0 +1,460 @@
|
||||
"""The seven gates: what each one rejects, in what order, and what it leaves on disk."""
|
||||
|
||||
import json
|
||||
import random
|
||||
from pathlib import Path
|
||||
|
||||
import nacl.bindings as sodium
|
||||
import pytest
|
||||
|
||||
from laborious.utils.bundle import format as fmt
|
||||
from laborious.utils.bundle.reader import BundleGateError, BundleLimits, BundleReader
|
||||
from laborious.utils.bundle.steps import ImportStep
|
||||
from tests.helpers import bundle_factory as factory
|
||||
|
||||
OBJECT_KEY = 'imported_models/factory_model_v2.sientia'
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def reader() -> BundleReader:
|
||||
return BundleReader()
|
||||
|
||||
|
||||
def _bundle_file(work_dir: Path, payload: bytes) -> Path:
|
||||
target = work_dir / 'download.sientia'
|
||||
target.write_bytes(payload)
|
||||
return target
|
||||
|
||||
|
||||
def _open(reader: BundleReader, work_dir: Path, payload: bytes, password: str | None = None):
|
||||
return reader.open(
|
||||
_bundle_file(work_dir, payload),
|
||||
password if password is not None else factory.DEFAULT_PASSWORD,
|
||||
work_dir / 'extract',
|
||||
work_dir,
|
||||
)
|
||||
|
||||
|
||||
def _extracted_files(root: Path) -> list[str]:
|
||||
if not root.exists():
|
||||
return []
|
||||
return sorted(str(p.relative_to(root)) for p in root.rglob('*') if p.is_file())
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- gate 1
|
||||
|
||||
|
||||
def test_gate_1_accepts_a_well_shaped_object(reader):
|
||||
reader.check_object_limits(size=1024, key=OBJECT_KEY)
|
||||
|
||||
|
||||
def test_gate_1_rejects_an_oversized_object(reader):
|
||||
limited = BundleReader(BundleLimits(max_object_bytes=1000))
|
||||
|
||||
with pytest.raises(BundleGateError) as error:
|
||||
limited.check_object_limits(size=1001, key=OBJECT_KEY)
|
||||
|
||||
assert error.value.gate == 1
|
||||
assert error.value.step is ImportStep.DOWNLOAD
|
||||
|
||||
|
||||
def test_gate_1_rejects_an_empty_object(reader):
|
||||
with pytest.raises(BundleGateError, match='empty'):
|
||||
reader.check_object_limits(size=0, key=OBJECT_KEY)
|
||||
|
||||
|
||||
def test_gate_1_rejects_a_key_outside_the_import_prefix(reader):
|
||||
with pytest.raises(BundleGateError, match='import prefix'):
|
||||
reader.check_object_limits(size=10, key='other_place/model.sientia')
|
||||
|
||||
|
||||
def test_gate_1_rejects_a_key_with_the_wrong_suffix(reader):
|
||||
with pytest.raises(BundleGateError, match='.sientia'):
|
||||
reader.check_object_limits(size=10, key='imported_models/model.zip')
|
||||
|
||||
|
||||
# ------------------------------------------------------------- golden bundle
|
||||
|
||||
|
||||
def test_opens_the_producer_fixture(golden_bundle, golden_password, work_dir, reader):
|
||||
"""The producer's own bytes, at the producer's own KDF cost."""
|
||||
opened = reader.open(golden_bundle, golden_password, work_dir / 'extract', work_dir)
|
||||
|
||||
assert opened.metadata['model_name'] == 'fixture_model'
|
||||
assert opened.metadata['experiment_name'] == 'fixture_experiment'
|
||||
assert set(opened.metadata) == {
|
||||
'model_name',
|
||||
'experiment_name',
|
||||
'run_id',
|
||||
'model_version',
|
||||
'model_project',
|
||||
'export_timestamp',
|
||||
}
|
||||
assert opened.parameters['target_variable'] == 'SE'
|
||||
assert opened.header is not None
|
||||
assert opened.digest == opened.header.ciphertext_digest.hex()
|
||||
assert _extracted_files(opened.extracted_dir) == [
|
||||
'artifacts/data_model/MLmodel',
|
||||
'artifacts/data_model/conda.yaml',
|
||||
'artifacts/data_model/model.pkl',
|
||||
'artifacts/data_model/python_env.yaml',
|
||||
'artifacts/data_model/requirements.txt',
|
||||
'artifacts/model_card/model_card.json',
|
||||
'artifacts/prediction_model/MLmodel',
|
||||
'artifacts/prediction_model/conda.yaml',
|
||||
'artifacts/prediction_model/model.pkl',
|
||||
'artifacts/prediction_model/python_env.yaml',
|
||||
'artifacts/prediction_model/requirements.txt',
|
||||
'artifacts/raw_data.csv',
|
||||
'artifacts/test_data.csv',
|
||||
'artifacts/train_data.csv',
|
||||
'metadata.json',
|
||||
]
|
||||
assert list(work_dir.glob('*.zip')) == []
|
||||
|
||||
|
||||
def test_fixture_metadata_matches_its_committed_plaintext_copy(
|
||||
golden_bundle, golden_password, work_dir, reader
|
||||
):
|
||||
from tests.laborious.utils.bundle.conftest import GOLDEN_METADATA
|
||||
|
||||
opened = reader.open(golden_bundle, golden_password, work_dir / 'extract', work_dir)
|
||||
committed = json.loads(GOLDEN_METADATA.read_text(encoding='utf-8'))
|
||||
|
||||
assert opened.metadata == committed['metadata']
|
||||
assert opened.parameters == committed['parameters']
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- gate 2
|
||||
|
||||
|
||||
def test_gate_2_rejects_a_wrong_magic_before_deriving(work_dir, reader, monkeypatch):
|
||||
payload = factory.make_bundle(magic=b'NOTSIENT')
|
||||
monkeypatch.setattr(
|
||||
sodium, 'crypto_pwhash_alg', _never_called('key derivation must not run at gate 2')
|
||||
)
|
||||
|
||||
with pytest.raises(BundleGateError) as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 2
|
||||
assert error.value.step is ImportStep.DOWNLOAD
|
||||
|
||||
|
||||
def test_gate_2_rejects_kdf_parameters_above_the_ceiling_without_allocating(
|
||||
work_dir, reader, monkeypatch
|
||||
):
|
||||
payload = factory.make_bundle(declared_memlimit=512 * 1024 * 1024)
|
||||
monkeypatch.setattr(
|
||||
sodium, 'crypto_pwhash_alg', _never_called('the ceiling must be enforced before deriving')
|
||||
)
|
||||
|
||||
with pytest.raises(BundleGateError) as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 2
|
||||
|
||||
|
||||
def test_gate_2_rejects_a_digest_mismatch(work_dir, reader):
|
||||
with pytest.raises(BundleGateError) as error:
|
||||
_open(reader, work_dir, factory.make_bundle(digest_override=bytes(32)))
|
||||
|
||||
assert error.value.gate == 2
|
||||
|
||||
|
||||
def test_gate_2_rejects_a_truncated_transfer(work_dir, reader):
|
||||
with pytest.raises(BundleGateError) as error:
|
||||
_open(reader, work_dir, factory.truncate_bundle(factory.make_bundle()))
|
||||
|
||||
assert error.value.gate == 2
|
||||
|
||||
|
||||
def _never_called(message: str):
|
||||
def _raise(*_args, **_kwargs):
|
||||
raise AssertionError(message)
|
||||
|
||||
return _raise
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- gate 3
|
||||
|
||||
AMBIGUOUS_REASON = 'wrong password, or the file is damaged or was altered'
|
||||
|
||||
|
||||
def test_gate_3_wrong_password(work_dir, reader):
|
||||
with pytest.raises(BundleGateError) as error:
|
||||
_open(reader, work_dir, factory.make_bundle(), password='Wrong1Password!Here')
|
||||
|
||||
assert error.value.gate == 3
|
||||
assert error.value.step is ImportStep.DECRYPTION
|
||||
assert AMBIGUOUS_REASON in str(error.value)
|
||||
|
||||
|
||||
def test_gate_3_altered_ciphertext_reads_the_same_as_a_wrong_password(work_dir, reader):
|
||||
tampered = factory.tamper_ciphertext(factory.make_bundle())
|
||||
|
||||
with pytest.raises(BundleGateError) as altered:
|
||||
_open(reader, work_dir, tampered)
|
||||
with pytest.raises(BundleGateError) as wrong_password:
|
||||
_open(reader, work_dir, factory.make_bundle(), password='Wrong1Password!Here')
|
||||
|
||||
assert str(altered.value) == str(wrong_password.value)
|
||||
assert altered.value.gate == wrong_password.value.gate == 3
|
||||
|
||||
|
||||
def test_gate_3_stream_without_a_final_tag_is_rejected(work_dir, reader):
|
||||
with pytest.raises(BundleGateError) as error:
|
||||
_open(reader, work_dir, factory.make_bundle(drop_final_tag=True))
|
||||
|
||||
assert error.value.gate == 3
|
||||
assert AMBIGUOUS_REASON in str(error.value)
|
||||
|
||||
|
||||
def test_gate_3_leaves_no_decrypted_zip_behind(work_dir, reader):
|
||||
with pytest.raises(BundleGateError):
|
||||
_open(reader, work_dir, factory.make_bundle(), password='Wrong1Password!Here')
|
||||
|
||||
assert list(work_dir.glob('*.zip')) == []
|
||||
|
||||
|
||||
def test_gate_3_normalises_the_password_to_nfc(work_dir, reader):
|
||||
composed = 'Café1Model!Test'
|
||||
decomposed = 'Café1Model!Test'
|
||||
payload = factory.make_bundle(password=composed)
|
||||
|
||||
opened = _open(reader, work_dir, payload, password=decomposed)
|
||||
|
||||
assert opened.metadata['model_name'] == 'factory_model'
|
||||
|
||||
|
||||
def test_gate_3_does_not_trim_the_password(work_dir, reader):
|
||||
payload = factory.make_bundle(password='Trailing1Space!Test ')
|
||||
|
||||
with pytest.raises(BundleGateError):
|
||||
_open(reader, work_dir, payload, password='Trailing1Space!Test')
|
||||
|
||||
|
||||
def test_multi_chunk_plaintext_round_trips(work_dir, reader):
|
||||
"""More than 64 KiB of plaintext, so `MESSAGE` and `FINAL` framing are both exercised."""
|
||||
spec = factory.ZipSpec()
|
||||
# Incompressible on purpose: a compressible filler would deflate below one chunk and the test
|
||||
# would silently stop covering multi-chunk framing.
|
||||
spec.files['artifacts/prediction_model/model.pkl'] = random.Random( # noqa: S311 - filler
|
||||
20260825
|
||||
).randbytes(200 * 1024)
|
||||
payload = factory.make_bundle(spec)
|
||||
assert len(payload) > fmt.CHUNK_SIZE + fmt.ABYTES + fmt.DIGEST_OFFSET
|
||||
|
||||
opened = _open(reader, work_dir, payload)
|
||||
|
||||
extracted = opened.extracted_dir / 'artifacts' / 'prediction_model' / 'model.pkl'
|
||||
assert extracted.read_bytes() == spec.files['artifacts/prediction_model/model.pkl']
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- gate 4
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('spec_kwargs', 'match'),
|
||||
[
|
||||
({'bomb_bytes': 40 * 1024 * 1024}, 'compression ratio'),
|
||||
({'traversal_entry': 'artifacts/../../etc/passwd'}, 'escaping'),
|
||||
({'absolute_entry': '/etc/passwd'}, 'absolute'),
|
||||
({'symlinks': {'artifacts/link': '/etc/passwd'}}, 'symbolic link'),
|
||||
],
|
||||
)
|
||||
def test_gate_4_rejects_hostile_central_directories(work_dir, reader, spec_kwargs, match):
|
||||
payload = factory.make_bundle(factory.ZipSpec(**spec_kwargs))
|
||||
|
||||
with pytest.raises(BundleGateError, match=match) as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 4
|
||||
assert error.value.step is ImportStep.ARCHIVE_INSPECTION
|
||||
assert _extracted_files(work_dir / 'extract') == []
|
||||
assert not (work_dir / 'extract').exists()
|
||||
|
||||
|
||||
def test_gate_4_rejects_too_many_entries(work_dir):
|
||||
limited = BundleReader(BundleLimits(max_entries=3))
|
||||
|
||||
with pytest.raises(BundleGateError, match='entries') as error:
|
||||
_open(limited, work_dir, factory.make_bundle())
|
||||
|
||||
assert error.value.gate == 4
|
||||
|
||||
|
||||
def test_gate_4_rejects_an_oversized_uncompressed_total(work_dir):
|
||||
limited = BundleReader(BundleLimits(max_uncompressed_bytes=64))
|
||||
|
||||
with pytest.raises(BundleGateError, match='uncompressed') as error:
|
||||
_open(limited, work_dir, factory.make_bundle())
|
||||
|
||||
assert error.value.gate == 4
|
||||
|
||||
|
||||
def test_gate_4_rejects_content_that_is_not_an_archive(work_dir, reader):
|
||||
payload = factory.encrypt_archive(b'not a zip at all')
|
||||
|
||||
with pytest.raises(BundleGateError, match='readable archive') as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 4
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- gate 5
|
||||
|
||||
|
||||
def test_gate_5_refuses_a_pre_existing_extraction_root(work_dir, reader):
|
||||
(work_dir / 'extract').mkdir()
|
||||
|
||||
with pytest.raises(BundleGateError, match='already exists') as error:
|
||||
_open(reader, work_dir, factory.make_bundle())
|
||||
|
||||
assert error.value.gate == 5
|
||||
assert error.value.step is ImportStep.EXTRACTION
|
||||
|
||||
|
||||
def test_gate_5_extraction_writes_plain_files(work_dir, reader):
|
||||
opened = _open(reader, work_dir, factory.make_bundle())
|
||||
|
||||
for path in opened.extracted_dir.rglob('*'):
|
||||
if path.is_file():
|
||||
assert path.stat().st_mode & 0o111 == 0, path
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- gate 6
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('spec', 'match'),
|
||||
[
|
||||
(factory.ZipSpec(include_metadata=False), 'metadata document'),
|
||||
(
|
||||
factory.ZipSpec(
|
||||
document=factory.build_metadata_document(drop_top_level_key='parameters')
|
||||
),
|
||||
'parameters and metadata blocks',
|
||||
),
|
||||
(
|
||||
factory.ZipSpec(
|
||||
document=factory.build_metadata_document(drop_metadata_field='experiment_name')
|
||||
),
|
||||
'experiment_name',
|
||||
),
|
||||
(
|
||||
factory.ZipSpec(
|
||||
document=factory.build_metadata_document(drop_parameter='target_variable')
|
||||
),
|
||||
'target_variable',
|
||||
),
|
||||
(
|
||||
factory.ZipSpec(
|
||||
document=factory.build_metadata_document(metadata={'experiment_name': 'x' * 51})
|
||||
),
|
||||
'experiment name',
|
||||
),
|
||||
(
|
||||
factory.ZipSpec(
|
||||
document=factory.build_metadata_document(metadata={'experiment_name': 'ab'})
|
||||
),
|
||||
'experiment name',
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_gate_6_rejects_a_malformed_metadata_document(work_dir, reader, spec, match):
|
||||
with pytest.raises(BundleGateError, match=match) as error:
|
||||
_open(reader, work_dir, factory.make_bundle(spec))
|
||||
|
||||
assert error.value.gate == 6
|
||||
assert error.value.step is ImportStep.STRUCTURE_VALIDATION
|
||||
assert not (work_dir / 'extract').exists()
|
||||
|
||||
|
||||
def test_gate_6_rejects_unreadable_metadata_json(work_dir, reader):
|
||||
spec = factory.ZipSpec(include_metadata=False)
|
||||
spec.extra_files['metadata.json'] = b'{not json'
|
||||
|
||||
with pytest.raises(BundleGateError, match='readable JSON') as error:
|
||||
_open(reader, work_dir, factory.make_bundle(spec))
|
||||
|
||||
assert error.value.gate == 6
|
||||
|
||||
|
||||
def test_gate_6_rejects_a_missing_artifacts_directory(work_dir, reader):
|
||||
spec = factory.ZipSpec(files={})
|
||||
|
||||
with pytest.raises(BundleGateError, match='artifacts directory') as error:
|
||||
_open(reader, work_dir, factory.make_bundle(spec))
|
||||
|
||||
assert error.value.gate == 6
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('dropped', 'match'),
|
||||
[
|
||||
('artifacts/data_model/MLmodel', 'data_model artifacts are incomplete'),
|
||||
('artifacts/prediction_model/model.pkl', 'prediction_model artifacts are incomplete'),
|
||||
],
|
||||
)
|
||||
def test_gate_6_rejects_an_incomplete_model_directory(work_dir, reader, dropped, match):
|
||||
spec = factory.ZipSpec()
|
||||
spec.files.pop(dropped)
|
||||
|
||||
with pytest.raises(BundleGateError, match=match) as error:
|
||||
_open(reader, work_dir, factory.make_bundle(spec))
|
||||
|
||||
assert error.value.gate == 6
|
||||
|
||||
|
||||
def test_gate_6_rejects_a_bundle_without_data_model(work_dir, reader):
|
||||
spec = factory.ZipSpec()
|
||||
for name in list(spec.files):
|
||||
if name.startswith('artifacts/data_model/'):
|
||||
spec.files.pop(name)
|
||||
|
||||
with pytest.raises(BundleGateError, match='no data_model artifacts') as error:
|
||||
_open(reader, work_dir, factory.make_bundle(spec))
|
||||
|
||||
assert error.value.gate == 6
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- gate 7
|
||||
|
||||
|
||||
@pytest.mark.parametrize('name', ['artifacts/run.sh', 'artifacts/tool.bin', 'artifacts/mod.so'])
|
||||
def test_gate_7_rejects_content_a_model_export_should_not_carry(work_dir, reader, name):
|
||||
spec = factory.ZipSpec()
|
||||
spec.extra_files[name] = b'#!/bin/sh\n'
|
||||
|
||||
with pytest.raises(BundleGateError, match='should not carry') as error:
|
||||
_open(reader, work_dir, factory.make_bundle(spec))
|
||||
|
||||
assert error.value.gate == 7
|
||||
assert error.value.step is ImportStep.CONTENT_POLICY
|
||||
assert not (work_dir / 'extract').exists()
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ ordering
|
||||
|
||||
|
||||
def test_the_earlier_gate_wins(work_dir, reader):
|
||||
"""A bundle that would fail gate 4 and gate 6 must report gate 4, having extracted nothing."""
|
||||
spec = factory.ZipSpec(
|
||||
traversal_entry='artifacts/../escape.csv',
|
||||
document=factory.build_metadata_document(drop_parameter='target_variable'),
|
||||
)
|
||||
|
||||
with pytest.raises(BundleGateError) as error:
|
||||
_open(reader, work_dir, factory.make_bundle(spec))
|
||||
|
||||
assert error.value.gate == 4
|
||||
assert _extracted_files(work_dir / 'extract') == []
|
||||
|
||||
|
||||
def test_header_is_checked_before_the_password(work_dir, reader):
|
||||
"""A bundle with a bad header and a wrong password reports the header, not the password."""
|
||||
with pytest.raises(BundleGateError) as error:
|
||||
_open(reader, work_dir, factory.make_bundle(magic=b'NOTSIENT'), password='Nope1!Nope1!')
|
||||
|
||||
assert error.value.gate == 2
|
||||
Reference in New Issue
Block a user