File size: 4,200 Bytes
db268d9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | """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
|