[DRAFT] TEST: Add notebook integration test for scenarios

pyrit/scenario/core/atomic_attack.py

@@ -23,7 +23,7 @@
     SingleTurnAttackContext,
 )
 from pyrit.executor.attack.core.attack_executor import AttackExecutorResult
-from pyrit.models import AttackResult, Message, SeedGroup
+from pyrit.models import AttackResult, Message
 
 logger = logging.getLogger(__name__)
 
@@ -93,7 +93,7 @@ def __init__(
         attack: AttackStrategy,
         objectives: List[str],
         prepended_conversations: Optional[List[List[Message]]] = None,
-        seed_groups: Optional[List[SeedGroup]] = None,
+        messages: Optional[List[Message]] = None,
         custom_prompts: Optional[List[str]] = None,
         memory_labels: Optional[Dict[str, str]] = None,
         **attack_execute_params: Any,
@@ -109,7 +109,7 @@ def __init__(
             prepended_conversations (Optional[List[List[Message]]]): Optional
                 list of conversation histories to prepend to each attack execution. This will be
                 used for all objectives.
-            seed_groups (Optional[List[SeedGroup]]): List of seed groups
+            messages (Optional[List[Message]]): List of messages
                 for single-turn attacks. Only valid for single-turn attacks.
             custom_prompts (Optional[List[str]]): List of custom prompts for multi-turn attacks.
                 Only valid for multi-turn attacks.
@@ -120,8 +120,6 @@ def __init__(
 
         Raises:
             ValueError: If objectives list is empty.
-            TypeError: If seed_groups is provided for multi-turn attacks or
-                custom_prompts is provided for single-turn attacks.
         """
         self.atomic_attack_name = atomic_attack_name
 
@@ -134,16 +132,9 @@ def __init__(
         # Determine context type once during initialization
         self._context_type: Literal["single_turn", "multi_turn", "unknown"] = self._determine_context_type(attack)
 
-        # Validate attack context type and parameters
-        self._validate_parameters(
-            seed_groups=seed_groups,
-            custom_prompts=custom_prompts,
-        )
-
         self._objectives = objectives
         self._prepended_conversations = prepended_conversations
-        self._seed_groups = seed_groups
-        self._custom_prompts = custom_prompts
+        self._messages = messages
         self._memory_labels = memory_labels or {}
         self._attack_execute_params = attack_execute_params
 
@@ -181,36 +172,6 @@ def _determine_context_type(self, attack: AttackStrategy) -> Literal["single_tur
                 return "multi_turn"
         return "unknown"
 
-    def _validate_parameters(
-        self,
-        *,
-        seed_groups: Optional[List[SeedGroup]],
-        custom_prompts: Optional[List[str]],
-    ) -> None:
-        """
-        Validate that parameters match the attack context type.
-
-        Args:
-            seed_groups (Optional[List[SeedGroup]]): Seed groups parameter.
-            custom_prompts (Optional[List[str]]): Custom prompts parameter.
-
-        Raises:
-            TypeError: If parameters don't match the attack context type.
-        """
-        # Validate seed_groups is only used with single-turn attacks
-        if seed_groups is not None and self._context_type != "single_turn":
-            raise TypeError(
-                f"seed_groups can only be used with single-turn attacks. "
-                f"Attack {self._attack.__class__.__name__} uses {self._context_type} context"
-            )
-
-        # Validate custom_prompts is only used with multi-turn attacks
-        if custom_prompts is not None and self._context_type != "multi_turn":
-            raise TypeError(
-                f"custom_prompts can only be used with multi-turn attacks. "
-                f"Attack {self._attack.__class__.__name__} uses {self._context_type} context"
-            )
-
     async def run_async(
         self,
         *,
@@ -268,7 +229,7 @@ async def run_async(
                 results = await executor.execute_single_turn_attacks_async(
                     attack=self._attack,
                     objectives=self._objectives,
-                    seed_groups=self._seed_groups,
+                    messages=self._messages,
                     prepended_conversations=prepended_conversations,
                     memory_labels=merged_memory_labels,
                     return_partial_on_failure=return_partial_on_failure,
@@ -278,7 +239,7 @@ async def run_async(
                 results = await executor.execute_multi_turn_attacks_async(
                     attack=self._attack,
                     objectives=self._objectives,
-                    custom_prompts=self._custom_prompts,
+                    messages=self._messages,
                     prepended_conversations=prepended_conversations,
                     memory_labels=merged_memory_labels,
                     return_partial_on_failure=return_partial_on_failure,

tests/integration/scenarios/test_notebooks_scenarios.py

@@ -0,0 +1,35 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import os
+import pathlib
+
+import nbformat
+import pytest
+from nbconvert.preprocessors import ExecutePreprocessor
+
+from pyrit.common import path
+
+nb_directory_path = pathlib.Path(path.DOCS_CODE_PATH, "scenarios").resolve()
+folder_paths = [d for d in nb_directory_path.iterdir() if d.is_dir()]
+
+
+@pytest.mark.parametrize(
+    "file_path",
+    [
+        os.path.join(dir_path, file)
+        for dir_path in folder_paths
+        for file in os.listdir(dir_path)
+        if file.endswith(".ipynb")
+    ],
+)
+def test_execute_notebooks(file_path):
+    nb_path = pathlib.Path(file_path).resolve()
+    print(nb_path)
+    with open(nb_path, encoding="utf-8") as f:
+        nb = nbformat.read(f, as_version=4)
+
+    ep = ExecutePreprocessor(timeout=900)
+
+    # Execute notebook, test will throw exception if any cell fails
+    ep.preprocess(nb, {"metadata": {"path": nb_path.parent}})