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