add basic tests for encoding probe

leondz · Jul 25, 2023 · 4491b56 · 4491b56
1 parent de7bac4
commit 4491b56
Showing 1 changed file with 28 additions and 0 deletions.
diff --git a/tests/probes/test_encoding.py b/tests/probes/test_encoding.py
@@ -0,0 +1,28 @@
+#!/usr/bin/env python3
+
+import garak.probes.encoding
+
+
+def test_InjectBase64_len_cap():
+    p = garak.probes.encoding.InjectBase64()
+    num_payloads = len(garak.probes.encoding.payloads)
+    num_templates = len(garak.probes.encoding.templates)
+    num_encoders = len(p.encoding_funcs)
+    assert len(p.prompts) < num_payloads * num_templates * num_encoders
+
+
+def test_InjectBase64_prompt_trigger_match():
+    p = garak.probes.encoding.InjectBase64()
+    num_payloads = len(garak.probes.encoding.payloads)
+    num_templates = len(garak.probes.encoding.templates)
+    num_encoders = len(p.encoding_funcs)
+    assert len(p.prompts) == len(p.triggers)
+
+
+def test_InjectBase64_triggers_not_in_prompts():
+    p = garak.probes.encoding.InjectBase64()
+    num_payloads = len(garak.probes.encoding.payloads)
+    num_templates = len(garak.probes.encoding.templates)
+    num_encoders = len(p.encoding_funcs)
+    for i in range(len(p.triggers)):
+        assert p.triggers[i] not in p.prompts[i]