tests/test_smoke.py

"""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)