FIX Reject num_templates and num_jailbreaks values below 1 in the jailbreak template sampler

pyrit/datasets/jailbreak/text_jailbreak.py

@@ -211,20 +211,24 @@ def get_jailbreak_templates(cls, num_templates: int | None = None) -> list[str]:
         Retrieve all jailbreaks from the JAILBREAK_TEMPLATES_PATH.
 
         Args:
-            num_templates (int, optional): Number of jailbreak templates to return. None to get all.
+            num_templates (int | None): Number of jailbreak templates to return. None to get all.
 
         Returns:
             list[str]: List of jailbreak template file names.
 
         Raises:
+            ValueError: If ``num_templates`` is provided and is less than 1.
             ValueError: If no jailbreak templates are found in the jailbreak directory.
             ValueError: If n is larger than the number of templates that exist.
         """
+        if num_templates is not None and num_templates < 1:
+            raise ValueError("num_templates must be at least 1")
+
         jailbreak_template_names = sorted(cls._get_template_cache().keys())
         if not jailbreak_template_names:
             raise ValueError("No jailbreak templates found in the jailbreak directory")
 
-        if num_templates:
+        if num_templates is not None:
             if num_templates > len(jailbreak_template_names):
                 raise ValueError(
                     f"Attempted to pull {num_templates} jailbreaks from a dataset"

pyrit/scenario/scenarios/airt/jailbreak.py

@@ -220,12 +220,16 @@ def set_params_from_args(self, *, args: dict[str, Any]) -> None:
             args (dict[str, Any]): Raw scenario run parameters.
 
         Raises:
-            ValueError: If ``num_jailbreak_attempts`` is less than one.
+            ValueError: If ``num_jailbreak_attempts`` is less than one, or if ``num_jailbreaks``
+                is provided and is less than one.
         """
         super().set_params_from_args(args=args)
         num_attempts = self.params["num_jailbreak_attempts"]
         if num_attempts < 1:
             raise ValueError("num_jailbreak_attempts must be at least 1")
+        num_jailbreaks = self.params.get("num_jailbreaks")
+        if num_jailbreaks is not None and num_jailbreaks < 1:
+            raise ValueError("num_jailbreaks must be at least 1")
 
     @apply_defaults
     def __init__(

tests/unit/datasets/test_jailbreak_text.py

@@ -89,6 +89,30 @@ def test_get_jailbreak_templates_includes_subdirectory_templates():
     assert len(templates) > top_level_count, "Subdirectory templates should be included in the listing"
 
 
+@pytest.mark.parametrize("num_templates", [1, 3])
+def test_get_jailbreak_templates_returns_requested_count(num_templates: int) -> None:
+    all_templates = TextJailBreak.get_jailbreak_templates()
+    templates = TextJailBreak.get_jailbreak_templates(num_templates=num_templates)
+    assert len(templates) == num_templates
+    assert set(templates).issubset(all_templates)
+
+
+@pytest.mark.parametrize("num_templates", [0, -1])
+def test_get_jailbreak_templates_rejects_invalid_num_templates(num_templates: int) -> None:
+    with pytest.raises(ValueError, match="num_templates must be at least 1"):
+        TextJailBreak.get_jailbreak_templates(num_templates=num_templates)
+
+
+def test_get_jailbreak_templates_none_returns_all() -> None:
+    assert TextJailBreak.get_jailbreak_templates(num_templates=None) == sorted(TextJailBreak._get_template_cache())
+
+
+def test_get_jailbreak_templates_rejects_count_above_population() -> None:
+    num_templates = len(TextJailBreak.get_jailbreak_templates()) + 1
+    with pytest.raises(ValueError, match=f"Attempted to pull {num_templates} jailbreaks"):
+        TextJailBreak.get_jailbreak_templates(num_templates=num_templates)
+
+
 def test_all_templates_render_without_syntax_errors(jailbreak_dir):
     """Test that all jailbreak templates can be successfully rendered with a test prompt."""
     yaml_files = [f for f in jailbreak_dir.rglob("*.yaml") if "multi_parameter" not in f.parts]

tests/unit/scenario/airt/test_jailbreak.py

@@ -208,6 +208,13 @@ def test_rejects_non_positive_num_jailbreak_attempts(self, mock_objective_scorer
         with pytest.raises(ValueError, match="num_jailbreak_attempts must be at least 1"):
             scenario.set_params_from_args(args={"num_jailbreak_attempts": num_attempts})
 
+    @pytest.mark.parametrize("num_jailbreaks", [0, -1])
+    def test_rejects_non_positive_num_jailbreaks(self, mock_objective_scorer, num_jailbreaks: int) -> None:
+        scenario = Jailbreak(objective_scorer=mock_objective_scorer)
+
+        with pytest.raises(ValueError, match="num_jailbreaks must be at least 1"):
+            scenario.set_params_from_args(args={"num_jailbreaks": num_jailbreaks})
+
     async def test_default_draws_random_template_sample(
         self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups
     ):