File size: 3,656 Bytes
cd07add | 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 | import unittest
from app.ml.dataset import build_dataset, label_name, LABEL_NAMES
class DatasetBuilderTests(unittest.TestCase):
def test_labels_are_the_documented_four(self):
self.assertEqual(LABEL_NAMES, ["NORMAL", "DISPOSABLE", "ROLE", "SPAM_TRAP"])
self.assertEqual(label_name(0), "NORMAL")
self.assertEqual(label_name(1), "DISPOSABLE")
self.assertEqual(label_name(2), "ROLE")
self.assertEqual(label_name(3), "SPAM_TRAP")
def test_blocklist_domains_become_disposable(self):
rows = build_dataset(
blocklist={"tempmail.example", "10minutemail.example"},
free={"gmail.com"},
role_prefixes={"admin"},
normal_domains={"acme-corp.test"},
normal_locals={"john.smith", "maria.garcia"},
per_class_limit=50,
)
by_label = {}
for r in rows:
by_label.setdefault(r["label"], []).append(r["text"])
self.assertIn(1, by_label, "blocklist domains must yield DISPOSABLE rows")
self.assertTrue(any("tempmail.example" in t for t in by_label[1]))
def test_role_prefixes_become_role_label(self):
rows = build_dataset(
blocklist=set(),
free=set(),
role_prefixes={"admin", "info"},
normal_domains={"acme-corp.test"},
normal_locals={"john.smith"},
per_class_limit=50,
)
role_rows = [r for r in rows if r["label"] == 2]
self.assertTrue(role_rows, "role prefixes must yield ROLE rows")
self.assertTrue(any(r["text"].startswith("admin") for r in role_rows))
def test_normal_class_never_contains_blocklist_domains(self):
rows = build_dataset(
blocklist={"tempmail.example"},
free=set(),
role_prefixes=set(),
normal_domains={"acme-corp.test", "northwind.test"},
normal_locals={"john.smith", "maria.garcia", "kenji.tanaka"},
per_class_limit=50,
)
normal_texts = [r["text"] for r in rows if r["label"] == 0]
self.assertTrue(normal_texts, "NORMAL class must be populated")
self.assertTrue(all("tempmail.example" not in t for t in normal_texts))
def test_per_class_limit_is_respected(self):
rows = build_dataset(
blocklist={f"d{i}.example" for i in range(100)},
free=set(),
role_prefixes=set(),
normal_domains={"acme.test"},
normal_locals={"john.smith"},
per_class_limit=25,
)
counts = {}
for r in rows:
counts[r["label"]] = counts.get(r["label"], 0) + 1
self.assertTrue(all(c <= 25 for c in counts.values()), counts)
def test_text_format_is_local_sep_domain(self):
rows = build_dataset(
blocklist=set(),
free=set(),
role_prefixes={"admin"},
normal_domains={"acme.test"},
normal_locals={"john.smith"},
per_class_limit=5,
)
self.assertTrue(rows)
for r in rows:
self.assertIn(" [SEP] ", r["text"])
def test_deterministic_with_same_seed(self):
kwargs = dict(
blocklist={"tempmail.example"},
free={"gmail.com"},
role_prefixes={"admin"},
normal_domains={"acme.test"},
normal_locals={"john.smith"},
per_class_limit=20,
seed=7,
)
a = build_dataset(**kwargs)
b = build_dataset(**kwargs)
self.assertEqual(a, b, "same seed must produce identical rows")
if __name__ == "__main__":
unittest.main() |