import tempfile import time import unittest from pathlib import Path from fgai.feedback import FeedbackStore, feedback_matches, select_match, suggest_generalizations class FeedbackTests(unittest.TestCase): def test_feedback_store_can_clear_by_status_or_all(self): with tempfile.TemporaryDirectory() as directory: store = FeedbackStore(str(Path(directory) / "feedback.json")) store.add({"status": "expected", "entity": "alice", "stream_id": "windows", "field": "username"}) store.add({"status": "confirmed", "entity": "bob", "stream_id": "vpn", "field": "username"}) expected = store.clear(status="expected") self.assertEqual(expected, {"removed": 1, "remaining": 1, "status": "expected"}) self.assertEqual(store.entries()[0]["status"], "confirmed") all_items = store.clear() self.assertEqual(all_items, {"removed": 1, "remaining": 0, "status": "all"}) self.assertEqual(store.entries(), []) def test_feedback_scope_defaults_and_ids_are_stable(self): with tempfile.TemporaryDirectory() as directory: store = FeedbackStore(str(Path(directory) / "feedback.json")) no_value = store.add({"status": "expected", "entity": "alice", "stream_id": "windows", "field": "username"}) self.assertEqual(no_value["scope"], "field") with_value = store.add({"status": "expected", "entity": "alice", "stream_id": "windows", "field": "username", "value": "bob"}) self.assertEqual(with_value["scope"], "value") self.assertEqual(len(store.entries()), 2) again = store.add({"status": "confirmed", "entity": "alice", "stream_id": "windows", "field": "username"}) self.assertEqual(again["id"], no_value["id"]) self.assertEqual(len(store.entries()), 2) self.assertEqual(next(item for item in store.entries() if item["id"] == no_value["id"]).get("status"), "confirmed") def test_feedback_broad_scope_requires_expiry(self): with tempfile.TemporaryDirectory() as directory: store = FeedbackStore(str(Path(directory) / "feedback.json")) with self.assertRaises(ValueError): store.add({"status": "expected", "scope": "entity", "entity": "alice"}) entry = store.add({"status": "expected", "scope": "entity", "entity": "alice", "expires_at": int(time.time()) + 3600}) self.assertEqual(entry["scope"], "entity") self.assertEqual(entry["stream_id"], "") def test_feedback_stream_scope_requires_stream_id(self): with tempfile.TemporaryDirectory() as directory: store = FeedbackStore(str(Path(directory) / "feedback.json")) with self.assertRaises(ValueError): store.add({"status": "expected", "scope": "stream", "expires_at": int(time.time()) + 3600}) entry = store.add({"status": "expected", "scope": "stream", "stream_id": "vpn", "expires_at": int(time.time()) + 3600}) self.assertEqual(entry["entity"], "") def test_feedback_entity_scope_requires_entity(self): with tempfile.TemporaryDirectory() as directory: store = FeedbackStore(str(Path(directory) / "feedback.json")) with self.assertRaises(ValueError): store.add({"status": "expected", "scope": "entity", "expires_at": int(time.time()) + 3600}) def test_feedback_revoke_removes_single_rule_by_id(self): with tempfile.TemporaryDirectory() as directory: store = FeedbackStore(str(Path(directory) / "feedback.json")) one = store.add({"status": "expected", "entity": "alice", "stream_id": "windows", "field": "username"}) store.add({"status": "confirmed", "entity": "bob", "stream_id": "vpn", "field": "username"}) result = store.revoke(one["id"]) self.assertEqual(result, {"removed": 1, "remaining": 1, "id": one["id"]}) self.assertEqual(len(store.entries()), 1) missing = store.revoke("does-not-exist") self.assertEqual(missing, {"removed": 0, "remaining": 1, "id": "does-not-exist"}) def test_add_many_writes_once_and_dedupes_by_id(self): with tempfile.TemporaryDirectory() as directory: store = FeedbackStore(str(Path(directory) / "feedback.json")) result = store.add_many([ {"status": "false_positive", "entity": "alice", "stream_id": "windows", "field": "username", "value": "bob"}, {"status": "confirmed", "entity": "carol", "stream_id": "vpn", "field": "username", "value": "dave"}, # Same (stream_id, entity, field, value) as the first item -> same id -> collapses to one. {"status": "expected", "entity": "alice", "stream_id": "windows", "field": "username", "value": "bob"}, ]) self.assertEqual(len(result["added"]), 2) self.assertEqual(result["errors"], []) self.assertEqual(result["remaining"], 2) entries = store.entries() self.assertEqual(len(entries), 2) collapsed = next(item for item in entries if item["entity"] == "alice") self.assertEqual(collapsed["status"], "expected") def test_add_many_reports_errors_without_aborting_the_batch(self): with tempfile.TemporaryDirectory() as directory: store = FeedbackStore(str(Path(directory) / "feedback.json")) result = store.add_many([ {"status": "false_positive", "entity": "alice", "stream_id": "windows", "field": "username", "value": "bob"}, {"status": "expected", "scope": "entity", "entity": "carol"}, # broad scope, no expiry -> invalid ]) self.assertEqual(len(result["added"]), 1) self.assertEqual(len(result["errors"]), 1) self.assertEqual(result["errors"][0]["index"], 1) self.assertEqual(len(store.entries()), 1) def test_feedback_matches_respects_scope_and_precedence(self): value_rule = {"scope": "value", "entity": "alice", "stream_id": "windows", "field": "username", "value": "10.0.0.1", "status": "confirmed"} field_rule = {"scope": "field", "entity": "alice", "stream_id": "windows", "field": "username", "status": "expected"} entity_rule = {"scope": "entity", "entity": "alice", "status": "expected"} stream_rule = {"scope": "stream", "stream_id": "windows", "status": "expected"} self.assertTrue(feedback_matches(stream_rule, entity="alice", stream_id="windows", field="username", value="10.0.0.1")) self.assertFalse(feedback_matches(stream_rule, entity="alice", stream_id="vpn", field="username", value="10.0.0.1")) self.assertTrue(feedback_matches(entity_rule, entity="alice", stream_id="vpn", field="hostname", value="x")) self.assertFalse(feedback_matches(field_rule, entity="alice", stream_id="windows", field="hostname", value="x")) self.assertTrue(feedback_matches(value_rule, entity="alice", stream_id="windows", field="username", value="10.0.0.1")) self.assertFalse(feedback_matches(value_rule, entity="alice", stream_id="windows", field="username", value="10.0.0.2")) candidates = [stream_rule, entity_rule, field_rule, value_rule] self.assertEqual(select_match(candidates), value_rule) self.assertEqual(select_match([stream_rule, entity_rule]), entity_rule) def test_suggest_generalizations_groups_narrow_value_rules_by_stream_entity_field_status(self): entries = [ {"scope": "value", "status": "false_positive", "entity": "alice", "stream_id": "windows", "field": "username", "value": "bob", "id": "1"}, {"scope": "value", "status": "false_positive", "entity": "alice", "stream_id": "windows", "field": "username", "value": "carol", "id": "2"}, {"scope": "value", "status": "false_positive", "entity": "alice", "stream_id": "windows", "field": "username", "value": "dave", "id": "3"}, {"scope": "value", "status": "confirmed", "entity": "alice", "stream_id": "windows", "field": "username", "value": "eve", "id": "4"}, ] suggestions = suggest_generalizations(entries, min_values=3) self.assertEqual(len(suggestions), 1) self.assertEqual(suggestions[0]["value_count"], 3) self.assertEqual(suggestions[0]["status"], "false_positive") self.assertEqual(set(suggestions[0]["entry_ids"]), {"1", "2", "3"}) def test_suggest_generalizations_respects_min_values_threshold(self): entries = [ {"scope": "value", "status": "expected", "entity": "a", "stream_id": "s", "field": "f", "value": "x", "id": "1"}, {"scope": "value", "status": "expected", "entity": "a", "stream_id": "s", "field": "f", "value": "y", "id": "2"}, ] self.assertEqual(suggest_generalizations(entries, min_values=3), []) def test_suggest_generalizations_ignores_broader_scope_entries(self): entries = [ {"scope": "field", "status": "false_positive", "entity": "a", "stream_id": "s", "field": "f", "id": "1"}, {"scope": "entity", "status": "expected", "entity": "a", "id": "2"}, ] self.assertEqual(suggest_generalizations(entries, min_values=2), []) def test_suggest_generalizations_keeps_different_entities_separate(self): entries = [ {"scope": "value", "status": "false_positive", "entity": "alice", "stream_id": "s", "field": "f", "value": "1", "id": "1"}, {"scope": "value", "status": "false_positive", "entity": "alice", "stream_id": "s", "field": "f", "value": "2", "id": "2"}, {"scope": "value", "status": "false_positive", "entity": "bob", "stream_id": "s", "field": "f", "value": "3", "id": "3"}, ] self.assertEqual(suggest_generalizations(entries, min_values=2), [ {"stream_id": "s", "entity": "alice", "field": "f", "status": "false_positive", "value_count": 2, "sample_values": ["1", "2"], "entry_ids": ["1", "2"]} ]) if __name__ == "__main__": unittest.main()