Tag probe datasets with OWASP LLM Top 10 categories in the scan UI

tests/integration/routes/test_probe.py

@@ -80,8 +80,10 @@ def test_data_config_endpoint():
     # Verify each item in response matches REGISTRY format
     for item in data:
         assert isinstance(item, dict)
-        # Add assertions for expected fields based on REGISTRY structure
-        # This will depend on what fields are defined in the REGISTRY items
+        # Every entry carries an OWASP LLM Top 10 tag for the UI badge
+        assert "owasp" in item
+        assert item["owasp"]["id"].startswith("LLM")
+        assert item["owasp"]["title"]
 
 
 def test_refusal_rate():

tests/unit/probe_data/test_owasp.py

@@ -0,0 +1,41 @@
+from agentic_security.probe_data import REGISTRY
+from agentic_security.probe_data.owasp import (
+    OWASP_LLM_TOP_10,
+    annotate,
+    classify,
+    owasp_tag,
+)
+
+
+def test_classify_defaults_to_prompt_injection():
+    assert classify({"dataset_name": "walledai/JailbreakBench"}) == "LLM01"
+    assert classify({"dataset_name": "deepset/prompt-injections"}) == "LLM01"
+    assert classify({}) == "LLM01"
+
+
+def test_classify_overrides():
+    assert classify({"dataset_name": "DataLeak"}) == "LLM06"
+    assert classify({"dataset_name": "Hallucination"}) == "LLM09"
+    assert classify({"dataset_name": "Malwaregen"}) == "LLM02"
+    assert classify({"dataset_name": "AgenticBackend"}) == "LLM08"
+
+
+def test_owasp_tag_shape():
+    tag = owasp_tag({"dataset_name": "Hallucination"})
+    assert tag["id"] == "LLM09"
+    assert tag["title"] == OWASP_LLM_TOP_10["LLM09"]
+    assert tag["url"].startswith("https://owasp.org/")
+
+
+def test_annotate_covers_every_registry_entry():
+    annotated = annotate(REGISTRY)
+    assert len(annotated) == len(REGISTRY)
+    for entry in annotated:
+        tag = entry["owasp"]
+        assert tag["id"] in OWASP_LLM_TOP_10
+        assert tag["title"] == OWASP_LLM_TOP_10[tag["id"]]
+
+
+def test_annotate_does_not_mutate_registry():
+    annotate(REGISTRY)
+    assert all("owasp" not in entry for entry in REGISTRY)