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