unit tests

Author: matthewreed26

src/inspect_evals/sec_qa/__init__.py

@@ -1,3 +1,19 @@
-from .sec_qa import sec_qa_v1, sec_qa_v1_5_shot, sec_qa_v2, sec_qa_v2_5_shot
+from .sec_qa import (
+    format_system_message_fewshot,
+    record_to_sample,
+    sample_to_fewshot,
+    sec_qa_v1,
+    sec_qa_v1_5_shot,
+    sec_qa_v2,
+    sec_qa_v2_5_shot,
+)
 
-__all__ = ["sec_qa_v1", "sec_qa_v1_5_shot", "sec_qa_v2", "sec_qa_v2_5_shot"]
+__all__ = [
+    "format_system_message_fewshot",
+    "record_to_sample",
+    "sample_to_fewshot",
+    "sec_qa_v1",
+    "sec_qa_v1_5_shot",
+    "sec_qa_v2",
+    "sec_qa_v2_5_shot",
+]

src/inspect_evals/sec_qa/sec_qa.py

@@ -45,13 +45,7 @@ def sec_qa_v1_5_shot() -> Task:
     return Task(
         dataset=dataset,
         solver=[
-            system_message(
-                SYSTEM_MESSAGE_FEWSHOT.format(
-                    examples="\n\n".join(
-                        [sample_to_fewshot(sample=sample) for sample in fewshot_samples]
-                    )
-                )
-            ),
+            system_message(format_system_message_fewshot(fewshot_samples)),
             multiple_choice(),
         ],
         scorer=choice(),
@@ -77,13 +71,7 @@ def sec_qa_v2_5_shot() -> Task:
     return Task(
         dataset=dataset,
         solver=[
-            system_message(
-                SYSTEM_MESSAGE_FEWSHOT.format(
-                    examples="\n\n".join(
-                        [sample_to_fewshot(sample=sample) for sample in fewshot_samples]
-                    )
-                )
-            ),
+            system_message(format_system_message_fewshot(fewshot_samples)),
             multiple_choice(),
         ],
         scorer=choice(),
@@ -108,13 +96,21 @@ def record_to_sample(record: dict[str, Any]) -> Sample:
     )
 
 
+def format_system_message_fewshot(fewshot_samples: Dataset) -> str:
+    return SYSTEM_MESSAGE_FEWSHOT.format(
+        examples="\n\n".join(
+            [sample_to_fewshot(sample=sample) for sample in fewshot_samples]
+        )
+    )
+
+
 def sample_to_fewshot(sample: Sample) -> str:
     prob_str = f"""QUESTION:\n{sample.input}"""
     labeled_choices = "\n"
     for i, letter_label in enumerate(["A", "B", "C", "D"]):
         labeled_choices = (
             labeled_choices
-            + f"{letter_label}: {sample.choices[i] if sample.choices is not None else None}\n"
+            + f"{letter_label}) {sample.choices[i] if sample.choices is not None else None}\n"
         )
     choices_str = f"""CHOICES:{labeled_choices}"""
     ans_str = f"""ANSWER: {sample.target}"""

tests/sec_qa/test_sec_qa.py

@@ -0,0 +1,53 @@
+from typing import Any
+
+from inspect_ai.dataset import Sample
+
+from inspect_evals.sec_qa import record_to_sample
+from inspect_evals.sec_qa.sec_qa import format_system_message_fewshot, sample_to_fewshot
+
+EXAMPLE_RECORD: dict[str, Any] = {
+    "Question": "Which of the following is a common indicator of an SQL injection attack?",
+    "A": "Frequent changes in user account permissions.",
+    "B": "Decreased performance of endpoint protection systems.",
+    "C": "Unusually high data upload rates to a web server.",
+    "D": "Sudden uptick in SQL queries, far beyond the usual baseline for the application.",
+    "Answer": "D",
+}
+
+EXPECTED_FORMATTED_SAMPLE_TO_FEWSHOT = "QUESTION:\nWhich of the following is a common indicator of an SQL injection attack?\n\nCHOICES:\nA) Frequent changes in user account permissions.\nB) Decreased performance of endpoint protection systems.\nC) Unusually high data upload rates to a web server.\nD) Sudden uptick in SQL queries, far beyond the usual baseline for the application.\n\n\nANSWER: D"
+
+
+def test_record_to_sample():
+    """Test that the record is mapped correctly to Inspect Sample"""
+    example_sample = record_to_sample(EXAMPLE_RECORD)
+    assert example_sample == Sample(
+        input="Which of the following is a common indicator of an SQL injection attack?",
+        choices=[
+            "Frequent changes in user account permissions.",
+            "Decreased performance of endpoint protection systems.",
+            "Unusually high data upload rates to a web server.",
+            "Sudden uptick in SQL queries, far beyond the usual baseline for the application.",
+        ],
+        target="D",
+    )
+
+
+def test_format_system_message_fewshot():
+    """Test that the system message is formatted correctly for some samples"""
+    formatted_system_message = format_system_message_fewshot(
+        [record_to_sample(EXAMPLE_RECORD), record_to_sample(EXAMPLE_RECORD)]
+    )
+    assert (
+        formatted_system_message
+        == "The following are multiple choice questions about Computer Security. Some examples are provided below."
+        + "\n\n"
+        + EXPECTED_FORMATTED_SAMPLE_TO_FEWSHOT
+        + "\n\n"
+        + EXPECTED_FORMATTED_SAMPLE_TO_FEWSHOT
+    )
+
+
+def test_sample_to_fewshot():
+    """Test that the sample is formatted correctly"""
+    formatted_sample_to_fewshot = sample_to_fewshot(record_to_sample(EXAMPLE_RECORD))
+    assert formatted_sample_to_fewshot == EXPECTED_FORMATTED_SAMPLE_TO_FEWSHOT