From f1054bff3cb167ba6c596dd312d5b6ffc6a727e6 Mon Sep 17 00:00:00 2001
From: sjadler2004 <sjadler2004@gmail.com>
Date: Tue, 17 Oct 2023 17:15:39 -0700
Subject: [PATCH] attempt to add tests for exclusions

---
 evals/elsuite/basic/includes_test.py | 25 +++++++++++++++++--------
 1 file changed, 17 insertions(+), 8 deletions(-)

diff --git a/evals/elsuite/basic/includes_test.py b/evals/elsuite/basic/includes_test.py
index f4a8c960f0..8d2f000aa9 100644
--- a/evals/elsuite/basic/includes_test.py
+++ b/evals/elsuite/basic/includes_test.py
@@ -11,13 +11,17 @@
 
 
 @mark.parametrize(
-    "completion, ideal, expected_match, ignore_case",
+    "completion, ideal, expected_match, ignore_case, use_exclusions, excluded_terms",
     [
-        ("world", "world", True, False),
-        ("world", "wOrLd", True, True),
-        ("world", ["world"], True, False),
-        ("world", ["foo", "bar"], False, False),
-        ("world", ["worldfoo", "worldbar"], False, False),
+        ("world", "world", True, False, False, None),
+        ("world", "wOrLd", True, True, False, None),
+        ("world", ["world"], True, False, False, None),
+        ("world", ["foo", "bar"], False, False, False, None),
+        ("world", ["worldfoo", "worldbar"], False, False, False, None),
+        # test for exclusions: does including an excludable word lead to False match, on what would otherwise be True match?
+        ("world exclusion", "world", False, False, True, ["exlusion", "excluded"]),
+        # change the 2nd word in completion from an excluded word, to make sure this would be true otherwise
+        ("world okay", "world", True, False, True, ["exlusion", "excluded"]),
     ],
 )
 def test_eval_sample(
@@ -25,21 +29,26 @@ def test_eval_sample(
     ideal: Union[str, list[str]],
     expected_match: bool,
     ignore_case: bool,
+    use_exclusions: bool = False,
+    excluded_terms: Union[str, list[str]] = []
 ):
     eval = Includes(
         completion_fns=[TestCompletionFn(completion)],
         samples_jsonl="",
         eval_registry_path=Path("."),
         ignore_case=ignore_case,
+        use_exclusions=use_exclusions,
+        excluded_terms=excluded_terms,
     )
 
     recorder = DummyRecorder(None)
     with recorder.as_default_recorder("x"), patch.object(
         recorder, "record_match", wraps=recorder.record_match
     ) as record_match:
-        eval.eval_sample(dict(input="Hello", ideal=ideal), None)
+        sample_dict = dict(input="Hello", ideal=ideal, exclude=excluded_terms)
+        eval.eval_sample(sample_dict, None)
         record_match.assert_called_once_with(
-            expected_match, expected=ideal, picked=completion, sampled=completion
+            expected_match, expected=ideal, picked=completion, sampled=completion,
         )