tests/test_data.py

"""Invariants the dataset must satisfy. These are the claims the cards make."""

import json
import sys
from collections import Counter
from pathlib import Path

import pytest

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from cti_attack import config, data  # noqa: E402

SCHEMES = ("document", "random")
SPLITS = ("train", "dev", "test")


def _built() -> bool:
    return (config.BUILD_DIR / "labels.json").exists()


pytestmark = pytest.mark.skipif(
    not _built(), reason="run scripts/01_build_dataset.py first")


@pytest.fixture(scope="module")
def labels():
    return data.load_labels()


@pytest.fixture(scope="module")
def splits():
    return {(s, p): data.load_split(s, p) for s in SCHEMES for p in SPLITS}


def test_no_document_spans_splits(splits):
    """The core claim of the document split: no report appears in two splits."""
    docs = {p: {r["doc_title"] for r in splits[("document", p)]} for p in SPLITS}
    assert not docs["train"] & docs["dev"]
    assert not docs["train"] & docs["test"]
    assert not docs["dev"] & docs["test"]


def test_random_split_does_leak_documents(splits):
    """The contrast case. If this ever passes cleanly the comparison is void."""
    docs = {p: {r["doc_title"] for r in splits[("random", p)]} for p in SPLITS}
    assert docs["train"] & docs["test"], "random split unexpectedly document-clean"


def test_every_technique_appears_in_every_document_split(splits, labels):
    for p in SPLITS:
        seen = {l for r in splits[("document", p)] for l in r["labels"]}
        missing = set(labels) - seen
        assert not missing, f"{p} is missing {sorted(missing)}"


def test_no_duplicate_sentences_within_scheme(splits):
    for scheme in SCHEMES:
        keys = [data.dedup_key(r["sentence"])
                for p in SPLITS for r in splits[(scheme, p)]]
        dupes = [k for k, c in Counter(keys).items() if c > 1]
        assert not dupes, f"{scheme}: {len(dupes)} duplicate sentences survived"


def test_dropped_technique_is_absent(splits, labels):
    stats = json.loads((config.BUILD_DIR / "build_stats.json").read_text(encoding="utf-8"))
    assert "T1557.001" in stats["dropped_techniques"]
    assert "T1557.001" not in labels
    for key, rows in splits.items():
        assert not any("T1557.001" in r["labels"] for r in rows), key


def test_label_count_matches_card(labels):
    assert len(labels) == 49


def test_boilerplate_is_stripped(splits):
    for rows in splits.values():
        assert not any(r["sentence"].lower().startswith("title:") for r in rows)


def test_clean_sentence_removes_header():
    raw = "title: Some Report url: https://example.com/x The dropper decodes its payload."
    assert data.clean_sentence(raw) == "The dropper decodes its payload."


def test_dataset_card_numbers_match_the_build(splits):
    """The dataset card must describe the dataset that actually exists.

    Added after the card was written from a pre-rebuild build and silently
    drifted by one or two rows on six different figures.
    """
    card = (config.REPO_ROOT / "DATASET_CARD.md").read_text(encoding="utf-8")
    stats = json.loads((config.BUILD_DIR / "build_stats.json").read_text(encoding="utf-8"))

    for key in ("final_sentences", "final_labelled", "duplicates_removed",
                "duplicate_groups", "cross_document_duplicates"):
        value = stats[key]
        assert f"{value:,}" in card or str(value) in card, \
            f"DATASET_CARD.md does not mention {key}={value}"

    for (scheme, split), rows in splits.items():
        n, labelled = len(rows), sum(1 for r in rows if r["labels"])
        assert f"{n:,}" in card, f"card missing {scheme}/{split} size {n:,}"
        assert f"{labelled:,}" in card, f"card missing {scheme}/{split} labelled {labelled:,}"


def test_dedup_merges_labels():
    stats = data.BuildStats()
    out = data.dedupe([
        {"sentence": "It decodes the payload.", "labels": ["T1027"], "doc_title": "a"},
        {"sentence": "It decodes the payload!", "labels": ["T1140"], "doc_title": "b"},
    ], stats)
    assert len(out) == 1
    assert out[0]["labels"] == ["T1027", "T1140"]
    assert stats.cross_document_duplicates == 1