"""End-to-end smoke test: load the model and run real threat prose through it. This is the test that stops the repo from shipping a checkpoint nobody ran. It does not assert accuracy — it asserts the artifact loads, produces well-formed output, and is wired to its thresholds and label set correctly. pytest tests/test_smoke.py -v Point it at a published model instead of a local one with: CTI_MODEL_DIR=your-username/cti-attack-mapper-modernbert pytest tests/test_smoke.py """ import os import sys from pathlib import Path import pytest sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) from cti_attack import config # noqa: E402 MODEL_DIR = os.environ.get( "CTI_MODEL_DIR", str(config.MODELS_DIR / "modernbert__document")) # Real sentences in the style of the training corpus, each with an ATT&CK # technique a analyst would expect to see fire. SAMPLES = [ "The dropper base64-encodes its configuration blob before writing it to disk.", "It then executes the payload through cmd.exe using a batch script.", "The implant establishes persistence by creating a scheduled task that runs at logon.", "Credentials were harvested from the LSASS process memory.", "The backdoor beacons to its command-and-control server over HTTPS.", ] def _available() -> bool: p = Path(MODEL_DIR) return p.exists() or "/" in MODEL_DIR pytestmark = pytest.mark.skipif( not _available(), reason=f"no model at {MODEL_DIR}; run scripts/03_train.py first") @pytest.fixture(scope="module") def mapper(): from cti_attack.predict import AttackMapper return AttackMapper(MODEL_DIR) def test_model_loads_with_expected_label_space(mapper): assert len(mapper.labels) == 49 assert all(l.startswith("T") for l in mapper.labels) assert len(mapper.thresholds) == len(mapper.labels) def test_single_prediction_is_well_formed(mapper): preds = mapper.predict(SAMPLES[0]) assert isinstance(preds, list) for p in preds: assert p.technique_id in mapper.labels assert 0.0 <= p.score <= 1.0 assert isinstance(p.name, str) and p.name def test_batch_matches_single(mapper): batch = mapper.predict_batch(SAMPLES) assert len(batch) == len(SAMPLES) single = mapper.predict(SAMPLES[2]) assert [p.technique_id for p in batch[2]] == [p.technique_id for p in single] def test_predictions_are_sorted_by_confidence(mapper): for preds in mapper.predict_batch(SAMPLES): scores = [p.score for p in preds] assert scores == sorted(scores, reverse=True) def test_model_fires_on_at_least_some_clear_cases(mapper): """Not an accuracy claim — a wiring check. If a model trained to macro-F1 ~0.5 predicts nothing at all across five textbook sentences, the thresholds or label mapping are broken. """ total = sum(len(p) for p in mapper.predict_batch(SAMPLES)) assert total > 0, "model produced no predictions on any sample sentence" def test_empty_and_odd_input_do_not_crash(mapper): assert mapper.predict_batch([]) == [] assert isinstance(mapper.predict(""), list) assert isinstance(mapper.predict("x" * 5000), list)