diff --git a/doc/code/scenarios/1_configuring_scenarios.ipynb b/doc/code/scenarios/1_configuring_scenarios.ipynb index 97cc8dbad..7a2b25958 100644 --- a/doc/code/scenarios/1_configuring_scenarios.ipynb +++ b/doc/code/scenarios/1_configuring_scenarios.ipynb @@ -36,8 +36,8 @@ "name": "stdout", "output_type": "stream", "text": [ - "Found default environment files: ['/home/vscode/.pyrit/.env']\n", - "Loaded environment file: /home/vscode/.pyrit/.env\n" + "Found default environment files: ['C:\\\\Users\\\\hannahwestra\\\\.pyrit\\\\.env']\n", + "Loaded environment file: C:\\Users\\hannahwestra\\.pyrit\\.env\n" ] } ], @@ -75,7 +75,7 @@ "output_type": "stream", "text": [ "\r", - "Loading datasets - this can take a few minutes: 0%| | 0/46 [00:00 AtomicAttack: + def _get_baseline(self) -> AtomicAttack: """ Get a baseline AtomicAttack, which simply sends all the objectives without any modifications. + If other atomic attacks exist, derives baseline data from the first attack. + Otherwise, creates a standalone baseline from the dataset configuration and scenario settings. + Returns: AtomicAttack: The baseline AtomicAttack instance. Raises: - ValueError: If no atomic attacks are available to derive baseline from. + ValueError: If required data (seed_groups, objective_target, attack_scoring_config) + is not available. """ - if not self._atomic_attacks or len(self._atomic_attacks) == 0: - raise ValueError("No atomic attacks available to derive baseline from.") - - first_attack = self._atomic_attacks[0] - - # Copy seed_groups, scoring, target from the first attack - seed_groups = first_attack.seed_groups - attack_scoring_config = first_attack._attack.get_attack_scoring_config() - objective_target = first_attack._attack.get_objective_target() - - if not seed_groups or len(seed_groups) == 0: - raise ValueError("First atomic attack must have seed_groups to create baseline.") - - if not objective_target: - raise ValueError("Objective target is required to create baseline attack.") - - if not attack_scoring_config: - raise ValueError("Attack scoring config is required to create baseline attack.") + seed_groups, attack_scoring_config, objective_target = self._get_baseline_data() # Create baseline attack with no converters attack = PromptSendingAttack( @@ -322,6 +315,40 @@ def _get_baseline_from_first_attack(self) -> AtomicAttack: memory_labels=self._memory_labels, ) + def _get_baseline_data(self) -> Tuple[List["SeedAttackGroup"], "AttackScoringConfig", PromptTarget]: + """ + Get the data needed to create a baseline attack. + + Returns the scenario-level data + + Returns: + Tuple containing (seed_groups, attack_scoring_config, objective_target) + + Raises: + ValueError: If required data is not available. + """ + # Create from scenario-level settings + if not self._objective_target: + raise ValueError("Objective target is required to create baseline attack.") + if not self._dataset_config: + raise ValueError("Dataset config is required to create baseline attack.") + if not self._objective_scorer: + raise ValueError("Objective scorer is required to create baseline attack.") + + seed_groups = self._dataset_config.get_all_seed_attack_groups() + if not seed_groups or len(seed_groups) == 0: + raise ValueError("Seed groups are required to create baseline attack.") + + # Import here to avoid circular imports + from pyrit.executor.attack.core.attack_config import AttackScoringConfig + + attack_scoring_config = AttackScoringConfig(objective_scorer=cast(TrueFalseScorer, self._objective_scorer)) + + if not attack_scoring_config: + raise ValueError("Attack scoring config is required to create baseline attack.") + + return seed_groups, attack_scoring_config, self._objective_target + def _raise_dataset_exception(self) -> None: error_msg = textwrap.dedent( f""" @@ -649,7 +676,8 @@ async def _execute_scenario_async(self) -> ScenarioResult: try: atomic_results = await atomic_attack.run_async( - max_concurrency=self._max_concurrency, return_partial_on_failure=True + max_concurrency=self._max_concurrency, + return_partial_on_failure=True, ) # Always save completed results, even if some objectives didn't complete @@ -676,7 +704,8 @@ async def _execute_scenario_async(self) -> ScenarioResult: # Mark scenario as failed self._memory.update_scenario_run_state( - scenario_result_id=scenario_result_id, scenario_run_state="FAILED" + scenario_result_id=scenario_result_id, + scenario_run_state="FAILED", ) # Raise exception with detailed information @@ -702,7 +731,8 @@ async def _execute_scenario_async(self) -> ScenarioResult: scenario_results = self._memory.get_scenario_results(scenario_result_ids=[scenario_result_id]) if scenario_results and scenario_results[0].scenario_run_state != "FAILED": self._memory.update_scenario_run_state( - scenario_result_id=scenario_result_id, scenario_run_state="FAILED" + scenario_result_id=scenario_result_id, + scenario_run_state="FAILED", ) raise diff --git a/pyrit/scenario/core/scenario_strategy.py b/pyrit/scenario/core/scenario_strategy.py index 362be2c56..d1f1cdceb 100644 --- a/pyrit/scenario/core/scenario_strategy.py +++ b/pyrit/scenario/core/scenario_strategy.py @@ -213,12 +213,14 @@ def prepare_scenario_strategies( strategies (Sequence[T | ScenarioCompositeStrategy] | None): The strategies to prepare. Can be a mix of bare strategy enums and composite strategies. If None, uses default_aggregate to determine defaults. + If an empty sequence, returns an empty list (useful for baseline-only execution). default_aggregate (T | None): The aggregate strategy to use when strategies is None. Common values: MyStrategy.ALL, MyStrategy.EASY. If None when strategies is None, raises ValueError. Returns: List[ScenarioCompositeStrategy]: Normalized list of composite strategies ready for use. + May be empty if an empty sequence was explicitly provided. Raises: ValueError: If strategies is None and default_aggregate is None, or if compositions @@ -251,7 +253,10 @@ def prepare_scenario_strategies( # For now, skip to allow flexibility pass + # Allow empty list if explicitly provided (for baseline-only execution) if not composite_strategies: + if strategies is not None and len(strategies) == 0: + return [] raise ValueError( f"No valid {cls.__name__} strategies provided. " f"Provide at least one {cls.__name__} enum or ScenarioCompositeStrategy." diff --git a/tests/unit/scenarios/test_cyber.py b/tests/unit/scenarios/test_cyber.py index 1730e7036..d84785190 100644 --- a/tests/unit/scenarios/test_cyber.py +++ b/tests/unit/scenarios/test_cyber.py @@ -15,6 +15,7 @@ from pyrit.identifiers import ScorerIdentifier from pyrit.models import SeedAttackGroup, SeedDataset, SeedObjective from pyrit.prompt_target import OpenAIChatTarget, PromptChatTarget, PromptTarget +from pyrit.scenario import DatasetConfiguration from pyrit.scenario.airt import Cyber, CyberStrategy from pyrit.score import TrueFalseCompositeScorer @@ -37,6 +38,16 @@ def mock_memory_seed_groups(): return [SeedAttackGroup(seeds=[SeedObjective(value=prompt)]) for prompt in seed_prompts] +@pytest.fixture +def mock_dataset_config(mock_memory_seed_groups): + """Create a mock dataset config that returns the seed groups.""" + mock_config = MagicMock(spec=DatasetConfiguration) + mock_config.get_all_seed_attack_groups.return_value = mock_memory_seed_groups + mock_config.get_default_dataset_names.return_value = ["airt_malware"] + mock_config.has_data_source.return_value = True + return mock_config + + @pytest.fixture def fast_cyberstrategy(): return CyberStrategy.SINGLE_TURN @@ -185,13 +196,13 @@ class TestCyberAttackGeneration: @pytest.mark.asyncio async def test_attack_generation_for_all( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, mock_dataset_config ): """Test that _get_atomic_attacks_async returns atomic attacks.""" with patch.object(Cyber, "_resolve_seed_groups", return_value=mock_memory_seed_groups): scenario = Cyber(objective_scorer=mock_objective_scorer) - await scenario.initialize_async(objective_target=mock_objective_target) + await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) atomic_attacks = await scenario._get_atomic_attacks_async() assert len(atomic_attacks) > 0 @@ -199,7 +210,12 @@ async def test_attack_generation_for_all( @pytest.mark.asyncio async def test_attack_generation_for_singleturn( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, fast_cyberstrategy + self, + mock_objective_target, + mock_objective_scorer, + mock_memory_seed_groups, + mock_dataset_config, + fast_cyberstrategy, ): """Test that the single turn attack generation works.""" with patch.object(Cyber, "_resolve_seed_groups", return_value=mock_memory_seed_groups): @@ -208,7 +224,9 @@ async def test_attack_generation_for_singleturn( ) await scenario.initialize_async( - objective_target=mock_objective_target, scenario_strategies=[fast_cyberstrategy] + objective_target=mock_objective_target, + scenario_strategies=[fast_cyberstrategy], + dataset_config=mock_dataset_config, ) atomic_attacks = await scenario._get_atomic_attacks_async() for run in atomic_attacks: @@ -216,7 +234,12 @@ async def test_attack_generation_for_singleturn( @pytest.mark.asyncio async def test_attack_generation_for_multiturn( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, slow_cyberstrategy + self, + mock_objective_target, + mock_objective_scorer, + mock_memory_seed_groups, + mock_dataset_config, + slow_cyberstrategy, ): """Test that the multi turn attack generation works.""" with patch.object(Cyber, "_resolve_seed_groups", return_value=mock_memory_seed_groups): @@ -225,7 +248,9 @@ async def test_attack_generation_for_multiturn( ) await scenario.initialize_async( - objective_target=mock_objective_target, scenario_strategies=[slow_cyberstrategy] + objective_target=mock_objective_target, + scenario_strategies=[slow_cyberstrategy], + dataset_config=mock_dataset_config, ) atomic_attacks = await scenario._get_atomic_attacks_async() @@ -234,7 +259,7 @@ async def test_attack_generation_for_multiturn( @pytest.mark.asyncio async def test_attack_runs_include_objectives( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, mock_dataset_config ): """Test that attack runs include objectives for each seed prompt.""" with patch.object(Cyber, "_resolve_seed_groups", return_value=mock_memory_seed_groups): @@ -242,7 +267,7 @@ async def test_attack_runs_include_objectives( objective_scorer=mock_objective_scorer, ) - await scenario.initialize_async(objective_target=mock_objective_target) + await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) atomic_attacks = await scenario._get_atomic_attacks_async() # Check that objectives are created for each seed prompt @@ -251,7 +276,7 @@ async def test_attack_runs_include_objectives( @pytest.mark.asyncio async def test_get_atomic_attacks_async_returns_attacks( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, mock_dataset_config ): """Test that _get_atomic_attacks_async returns atomic attacks.""" with patch.object(Cyber, "_resolve_seed_groups", return_value=mock_memory_seed_groups): @@ -259,7 +284,7 @@ async def test_get_atomic_attacks_async_returns_attacks( objective_scorer=mock_objective_scorer, ) - await scenario.initialize_async(objective_target=mock_objective_target) + await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) atomic_attacks = await scenario._get_atomic_attacks_async() assert len(atomic_attacks) > 0 assert all(hasattr(run, "_attack") for run in atomic_attacks) @@ -273,17 +298,19 @@ class TestCyberLifecycle: @pytest.mark.asyncio async def test_initialize_async_with_max_concurrency( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, mock_dataset_config ): """Test initialization with custom max_concurrency.""" with patch.object(Cyber, "_resolve_seed_groups", return_value=mock_memory_seed_groups): scenario = Cyber(objective_scorer=mock_objective_scorer) - await scenario.initialize_async(objective_target=mock_objective_target, max_concurrency=20) + await scenario.initialize_async( + objective_target=mock_objective_target, max_concurrency=20, dataset_config=mock_dataset_config + ) assert scenario._max_concurrency == 20 @pytest.mark.asyncio async def test_initialize_async_with_memory_labels( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, mock_dataset_config ): """Test initialization with memory labels.""" memory_labels = {"test": "cyber", "category": "scenario"} @@ -295,6 +322,7 @@ async def test_initialize_async_with_memory_labels( await scenario.initialize_async( memory_labels=memory_labels, objective_target=mock_objective_target, + dataset_config=mock_dataset_config, ) assert scenario._memory_labels == memory_labels @@ -316,11 +344,11 @@ def test_scenario_version_is_set(self, mock_objective_scorer, mock_memory_seed_g assert scenario.version == 1 @pytest.mark.asyncio - async def test_no_target_duplication(self, mock_objective_target, mock_memory_seed_groups): + async def test_no_target_duplication(self, mock_objective_target, mock_memory_seed_groups, mock_dataset_config): """Test that all three targets (adversarial, object, scorer) are distinct.""" with patch.object(Cyber, "_resolve_seed_groups", return_value=mock_memory_seed_groups): scenario = Cyber() - await scenario.initialize_async(objective_target=mock_objective_target) + await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) objective_target = scenario._objective_target diff --git a/tests/unit/scenarios/test_encoding.py b/tests/unit/scenarios/test_encoding.py index a3624029f..a980e43fc 100644 --- a/tests/unit/scenarios/test_encoding.py +++ b/tests/unit/scenarios/test_encoding.py @@ -9,9 +9,10 @@ from pyrit.executor.attack import PromptSendingAttack from pyrit.identifiers import ScorerIdentifier -from pyrit.models import SeedPrompt +from pyrit.models import SeedAttackGroup, SeedObjective, SeedPrompt from pyrit.prompt_converter import Base64Converter from pyrit.prompt_target import PromptTarget +from pyrit.scenario import DatasetConfiguration from pyrit.scenario.garak import Encoding, EncodingStrategy from pyrit.score import DecodingScorer, TrueFalseScorer @@ -37,6 +38,17 @@ def mock_memory_seeds(): ] +@pytest.fixture +def mock_dataset_config(mock_memory_seeds): + """Create a mock dataset config that returns the seed groups.""" + seed_groups = [SeedAttackGroup(seeds=[SeedObjective(value=seed.value)]) for seed in mock_memory_seeds] + mock_config = MagicMock(spec=DatasetConfiguration) + mock_config.get_all_seed_attack_groups.return_value = seed_groups + mock_config.get_default_dataset_names.return_value = ["garak_encoding"] + mock_config.has_data_source.return_value = True + return mock_config + + @pytest.fixture def mock_objective_target(): """Create a mock objective target for testing.""" @@ -158,7 +170,9 @@ def test_init_with_max_concurrency(self, mock_objective_target, mock_objective_s assert scenario._max_concurrency == 1 @pytest.mark.asyncio - async def test_init_attack_strategies(self, mock_objective_target, mock_objective_scorer, mock_memory_seeds): + async def test_init_attack_strategies( + self, mock_objective_target, mock_objective_scorer, mock_memory_seeds, mock_dataset_config + ): """Test that attack strategies are set correctly.""" from unittest.mock import patch @@ -167,7 +181,7 @@ async def test_init_attack_strategies(self, mock_objective_target, mock_objectiv objective_scorer=mock_objective_scorer, ) - await scenario.initialize_async(objective_target=mock_objective_target) + await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) # By default, EncodingStrategy.ALL is used, which expands to all encoding strategies assert len(scenario._scenario_composites) > 0 @@ -189,7 +203,7 @@ class TestEncodingAtomicAttacks: @pytest.mark.asyncio async def test_get_atomic_attacks_async_returns_attacks( - self, mock_objective_target, mock_objective_scorer, mock_memory_seeds + self, mock_objective_target, mock_objective_scorer, mock_memory_seeds, mock_dataset_config ): """Test that _get_atomic_attacks_async returns atomic attacks.""" from unittest.mock import patch @@ -199,7 +213,7 @@ async def test_get_atomic_attacks_async_returns_attacks( objective_scorer=mock_objective_scorer, ) - await scenario.initialize_async(objective_target=mock_objective_target) + await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) atomic_attacks = await scenario._get_atomic_attacks_async() # Should return multiple atomic attacks (one for each encoding type) @@ -208,7 +222,7 @@ async def test_get_atomic_attacks_async_returns_attacks( @pytest.mark.asyncio async def test_get_converter_attacks_returns_multiple_encodings( - self, mock_objective_target, mock_objective_scorer, mock_memory_seeds + self, mock_objective_target, mock_objective_scorer, mock_memory_seeds, mock_dataset_config ): """Test that _get_converter_attacks returns attacks for multiple encoding types.""" from unittest.mock import patch @@ -218,7 +232,7 @@ async def test_get_converter_attacks_returns_multiple_encodings( objective_scorer=mock_objective_scorer, ) - await scenario.initialize_async(objective_target=mock_objective_target) + await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) attack_runs = scenario._get_converter_attacks() # Should have multiple attack runs for different encodings @@ -228,7 +242,7 @@ async def test_get_converter_attacks_returns_multiple_encodings( @pytest.mark.asyncio async def test_get_prompt_attacks_creates_attack_runs( - self, mock_objective_target, mock_objective_scorer, mock_memory_seeds + self, mock_objective_target, mock_objective_scorer, mock_memory_seeds, mock_dataset_config ): """Test that _get_prompt_attacks creates attack runs with correct structure.""" from unittest.mock import patch @@ -238,7 +252,7 @@ async def test_get_prompt_attacks_creates_attack_runs( objective_scorer=mock_objective_scorer, ) - await scenario.initialize_async(objective_target=mock_objective_target) + await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) attack_runs = scenario._get_prompt_attacks(converters=[Base64Converter()], encoding_name="Base64") # Should create attack runs @@ -253,7 +267,7 @@ async def test_get_prompt_attacks_creates_attack_runs( @pytest.mark.asyncio async def test_attack_runs_include_objectives( - self, mock_objective_target, mock_objective_scorer, mock_memory_seeds + self, mock_objective_target, mock_objective_scorer, mock_memory_seeds, mock_dataset_config ): """Test that attack runs include objectives for each seed prompt.""" from unittest.mock import patch @@ -263,7 +277,7 @@ async def test_attack_runs_include_objectives( objective_scorer=mock_objective_scorer, ) - await scenario.initialize_async(objective_target=mock_objective_target) + await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) attack_runs = scenario._get_prompt_attacks(converters=[Base64Converter()], encoding_name="Base64") # Check that objectives are created for each seed prompt @@ -279,7 +293,9 @@ class TestEncodingExecution: """Tests for Encoding execution.""" @pytest.mark.asyncio - async def test_scenario_initialization(self, mock_objective_target, mock_objective_scorer, mock_memory_seeds): + async def test_scenario_initialization( + self, mock_objective_target, mock_objective_scorer, mock_memory_seeds, mock_dataset_config + ): """Test that scenario can be initialized successfully.""" from unittest.mock import patch @@ -288,14 +304,14 @@ async def test_scenario_initialization(self, mock_objective_target, mock_objecti objective_scorer=mock_objective_scorer, ) - await scenario.initialize_async(objective_target=mock_objective_target) + await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) # Verify initialization creates atomic attacks assert scenario.atomic_attack_count > 0 @pytest.mark.asyncio async def test_resolve_seed_prompts_loads_garak_data( - self, mock_objective_target, mock_objective_scorer, mock_memory_seeds + self, mock_objective_target, mock_objective_scorer, mock_memory_seeds, mock_dataset_config ): """Test that _resolve_seed_prompts loads data from Garak datasets.""" from unittest.mock import patch diff --git a/tests/unit/scenarios/test_foundry.py b/tests/unit/scenarios/test_foundry.py index 9dc960fcc..aef8fde09 100644 --- a/tests/unit/scenarios/test_foundry.py +++ b/tests/unit/scenarios/test_foundry.py @@ -15,7 +15,7 @@ from pyrit.prompt_converter import Base64Converter from pyrit.prompt_target import PromptTarget from pyrit.prompt_target.common.prompt_chat_target import PromptChatTarget -from pyrit.scenario import AtomicAttack +from pyrit.scenario import AtomicAttack, DatasetConfiguration from pyrit.scenario.foundry import FoundryStrategy, RedTeamAgent from pyrit.score import FloatScaleThresholdScorer, TrueFalseScorer @@ -42,6 +42,16 @@ def mock_memory_seed_groups(): return [SeedAttackGroup(seeds=[SeedObjective(value=obj)]) for obj in objectives] +@pytest.fixture +def mock_dataset_config(mock_memory_seed_groups): + """Create a mock dataset config that returns the seed groups.""" + mock_config = MagicMock(spec=DatasetConfiguration) + mock_config.get_all_seed_attack_groups.return_value = mock_memory_seed_groups + mock_config.get_default_dataset_names.return_value = ["foundry_red_team"] + mock_config.has_data_source.return_value = True + return mock_config + + @pytest.fixture def mock_objective_target(): """Create a mock objective target for testing.""" @@ -95,7 +105,7 @@ class TestFoundryInitialization: ) @pytest.mark.asyncio async def test_init_with_single_strategy( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, mock_dataset_config ): """Test initialization with a single attack strategy.""" with patch.object(RedTeamAgent, "_resolve_seed_groups", return_value=mock_memory_seed_groups): @@ -106,6 +116,7 @@ async def test_init_with_single_strategy( await scenario.initialize_async( objective_target=mock_objective_target, scenario_strategies=[FoundryStrategy.Base64], + dataset_config=mock_dataset_config, ) assert scenario.atomic_attack_count > 0 assert scenario.name == "RedTeamAgent" @@ -120,7 +131,7 @@ async def test_init_with_single_strategy( ) @pytest.mark.asyncio async def test_init_with_multiple_strategies( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, mock_dataset_config ): """Test initialization with multiple attack strategies.""" strategies = [ @@ -137,6 +148,7 @@ async def test_init_with_multiple_strategies( await scenario.initialize_async( objective_target=mock_objective_target, scenario_strategies=strategies, + dataset_config=mock_dataset_config, ) assert scenario.atomic_attack_count >= len(strategies) @@ -202,7 +214,9 @@ def test_init_with_custom_scorer(self, mock_objective_target, mock_objective_sco }, ) @pytest.mark.asyncio - async def test_init_with_memory_labels(self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups): + async def test_init_with_memory_labels( + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, mock_dataset_config + ): """Test initialization with memory labels.""" memory_labels = {"test": "foundry", "category": "attack"} @@ -216,6 +230,7 @@ async def test_init_with_memory_labels(self, mock_objective_target, mock_objecti await scenario.initialize_async( objective_target=mock_objective_target, memory_labels=memory_labels, + dataset_config=mock_dataset_config, ) assert scenario._memory_labels == memory_labels @@ -281,7 +296,7 @@ class TestFoundryStrategyNormalization: ) @pytest.mark.asyncio async def test_normalize_easy_strategies( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, mock_dataset_config ): """Test that EASY strategy expands to easy attack strategies.""" with patch.object(RedTeamAgent, "_resolve_seed_groups", return_value=mock_memory_seed_groups): @@ -292,6 +307,7 @@ async def test_normalize_easy_strategies( await scenario.initialize_async( objective_target=mock_objective_target, scenario_strategies=[FoundryStrategy.EASY], + dataset_config=mock_dataset_config, ) # EASY should expand to multiple attack strategies assert scenario.atomic_attack_count > 1 @@ -306,7 +322,7 @@ async def test_normalize_easy_strategies( ) @pytest.mark.asyncio async def test_normalize_moderate_strategies( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, mock_dataset_config ): """Test that MODERATE strategy expands to moderate attack strategies.""" with patch.object(RedTeamAgent, "_resolve_seed_groups", return_value=mock_memory_seed_groups): @@ -317,6 +333,7 @@ async def test_normalize_moderate_strategies( await scenario.initialize_async( objective_target=mock_objective_target, scenario_strategies=[FoundryStrategy.MODERATE], + dataset_config=mock_dataset_config, ) # MODERATE should expand to moderate attack strategies (currently only 1: Tense) assert scenario.atomic_attack_count >= 1 @@ -331,7 +348,7 @@ async def test_normalize_moderate_strategies( ) @pytest.mark.asyncio async def test_normalize_difficult_strategies( - self, mock_objective_target, mock_float_threshold_scorer, mock_memory_seed_groups + self, mock_objective_target, mock_float_threshold_scorer, mock_memory_seed_groups, mock_dataset_config ): """Test that DIFFICULT strategy expands to difficult attack strategies.""" with patch.object(RedTeamAgent, "_resolve_seed_groups", return_value=mock_memory_seed_groups): @@ -343,6 +360,7 @@ async def test_normalize_difficult_strategies( await scenario.initialize_async( objective_target=mock_objective_target, scenario_strategies=[FoundryStrategy.DIFFICULT], + dataset_config=mock_dataset_config, ) # DIFFICULT should expand to multiple attack strategies assert scenario.atomic_attack_count > 1 @@ -357,7 +375,7 @@ async def test_normalize_difficult_strategies( ) @pytest.mark.asyncio async def test_normalize_mixed_difficulty_levels( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, mock_dataset_config ): """Test that multiple difficulty levels expand correctly.""" with patch.object(RedTeamAgent, "_resolve_seed_groups", return_value=mock_memory_seed_groups): @@ -368,6 +386,7 @@ async def test_normalize_mixed_difficulty_levels( await scenario.initialize_async( objective_target=mock_objective_target, scenario_strategies=[FoundryStrategy.EASY, FoundryStrategy.MODERATE], + dataset_config=mock_dataset_config, ) # Combined difficulty levels should expand to multiple strategies assert scenario.atomic_attack_count > 5 # EASY has 20, MODERATE has 1, combined should have more @@ -382,7 +401,7 @@ async def test_normalize_mixed_difficulty_levels( ) @pytest.mark.asyncio async def test_normalize_with_specific_and_difficulty_levels( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, mock_dataset_config ): """Test that specific strategies combined with difficulty levels work correctly.""" with patch.object(RedTeamAgent, "_resolve_seed_groups", return_value=mock_memory_seed_groups): @@ -396,6 +415,7 @@ async def test_normalize_with_specific_and_difficulty_levels( FoundryStrategy.EASY, FoundryStrategy.Base64, # Specific strategy ], + dataset_config=mock_dataset_config, ) # EASY expands to 20 strategies, but Base64 might already be in EASY, so at least 20 assert scenario.atomic_attack_count >= 20 @@ -415,7 +435,7 @@ class TestFoundryAttackCreation: ) @pytest.mark.asyncio async def test_get_attack_from_single_turn_strategy( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, mock_dataset_config ): """Test creating an attack from a single-turn strategy.""" with patch.object(RedTeamAgent, "_resolve_seed_groups", return_value=mock_memory_seed_groups): @@ -426,6 +446,7 @@ async def test_get_attack_from_single_turn_strategy( await scenario.initialize_async( objective_target=mock_objective_target, scenario_strategies=[FoundryStrategy.Base64], + dataset_config=mock_dataset_config, ) # Get the composite strategy that was created during initialization @@ -445,7 +466,12 @@ async def test_get_attack_from_single_turn_strategy( ) @pytest.mark.asyncio async def test_get_attack_from_multi_turn_strategy( - self, mock_objective_target, mock_adversarial_target, mock_objective_scorer, mock_memory_seed_groups + self, + mock_objective_target, + mock_adversarial_target, + mock_objective_scorer, + mock_memory_seed_groups, + mock_dataset_config, ): """Test creating a multi-turn attack strategy.""" with patch.object(RedTeamAgent, "_resolve_seed_groups", return_value=mock_memory_seed_groups): @@ -457,6 +483,7 @@ async def test_get_attack_from_multi_turn_strategy( await scenario.initialize_async( objective_target=mock_objective_target, scenario_strategies=[FoundryStrategy.Crescendo], + dataset_config=mock_dataset_config, ) # Get the composite strategy that was created during initialization @@ -481,7 +508,7 @@ class TestFoundryGetAttack: ) @pytest.mark.asyncio async def test_get_attack_single_turn_with_converters( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, mock_dataset_config ): """Test creating a single-turn attack with converters.""" with patch.object(RedTeamAgent, "_resolve_seed_groups", return_value=mock_memory_seed_groups): @@ -492,6 +519,7 @@ async def test_get_attack_single_turn_with_converters( await scenario.initialize_async( objective_target=mock_objective_target, scenario_strategies=[FoundryStrategy.Base64], + dataset_config=mock_dataset_config, ) attack = scenario._get_attack( @@ -511,7 +539,12 @@ async def test_get_attack_single_turn_with_converters( ) @pytest.mark.asyncio async def test_get_attack_multi_turn_with_adversarial_target( - self, mock_objective_target, mock_adversarial_target, mock_objective_scorer, mock_memory_seed_groups + self, + mock_objective_target, + mock_adversarial_target, + mock_objective_scorer, + mock_memory_seed_groups, + mock_dataset_config, ): """Test creating a multi-turn attack.""" with patch.object(RedTeamAgent, "_resolve_seed_groups", return_value=mock_memory_seed_groups): @@ -523,6 +556,7 @@ async def test_get_attack_multi_turn_with_adversarial_target( await scenario.initialize_async( objective_target=mock_objective_target, scenario_strategies=[FoundryStrategy.Crescendo], + dataset_config=mock_dataset_config, ) attack = scenario._get_attack( @@ -573,7 +607,7 @@ class TestFoundryAllStrategies: ) @pytest.mark.asyncio async def test_all_single_turn_strategies_create_attack_runs( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, strategy + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, mock_dataset_config, strategy ): """Test that all single-turn strategies can create attack runs.""" with patch.object(RedTeamAgent, "_resolve_seed_groups", return_value=mock_memory_seed_groups): @@ -584,6 +618,7 @@ async def test_all_single_turn_strategies_create_attack_runs( await scenario.initialize_async( objective_target=mock_objective_target, scenario_strategies=[strategy], + dataset_config=mock_dataset_config, ) # Get the composite strategy that was created during initialization @@ -613,6 +648,7 @@ async def test_all_multi_turn_strategies_create_attack_runs( mock_adversarial_target, mock_objective_scorer, mock_memory_seed_groups, + mock_dataset_config, strategy, ): """Test that all multi-turn strategies can create attack runs.""" @@ -625,6 +661,7 @@ async def test_all_multi_turn_strategies_create_attack_runs( await scenario.initialize_async( objective_target=mock_objective_target, scenario_strategies=[strategy], + dataset_config=mock_dataset_config, ) # Get the composite strategy that was created during initialization @@ -647,7 +684,7 @@ class TestFoundryProperties: ) @pytest.mark.asyncio async def test_scenario_composites_set_after_initialize( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, mock_dataset_config ): """Test that scenario composites are set after initialize_async.""" strategies = [FoundryStrategy.Base64, FoundryStrategy.ROT13] @@ -664,6 +701,7 @@ async def test_scenario_composites_set_after_initialize( await scenario.initialize_async( objective_target=mock_objective_target, scenario_strategies=strategies, + dataset_config=mock_dataset_config, ) # After initialize_async, composites should be set @@ -696,7 +734,7 @@ def test_scenario_version_is_set(self, mock_objective_target, mock_objective_sco ) @pytest.mark.asyncio async def test_scenario_atomic_attack_count_matches_strategies( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, mock_dataset_config ): """Test that atomic attack count is reasonable for the number of strategies.""" strategies = [ @@ -713,6 +751,7 @@ async def test_scenario_atomic_attack_count_matches_strategies( await scenario.initialize_async( objective_target=mock_objective_target, scenario_strategies=strategies, + dataset_config=mock_dataset_config, ) # Should have at least as many runs as specific strategies provided assert scenario.atomic_attack_count >= len(strategies) diff --git a/tests/unit/scenarios/test_leakage_scenario.py b/tests/unit/scenarios/test_leakage_scenario.py index 3de049232..1d795a7e2 100644 --- a/tests/unit/scenarios/test_leakage_scenario.py +++ b/tests/unit/scenarios/test_leakage_scenario.py @@ -13,8 +13,9 @@ from pyrit.executor.attack import CrescendoAttack, PromptSendingAttack, RolePlayAttack from pyrit.executor.attack.core.attack_config import AttackScoringConfig from pyrit.identifiers import ScorerIdentifier -from pyrit.models import SeedDataset, SeedObjective +from pyrit.models import SeedAttackGroup, SeedDataset, SeedObjective from pyrit.prompt_target import OpenAIChatTarget, PromptChatTarget, PromptTarget +from pyrit.scenario import DatasetConfiguration from pyrit.scenario.airt import LeakageScenario, LeakageStrategy from pyrit.score import TrueFalseCompositeScorer @@ -36,6 +37,17 @@ def mock_memory_seeds(): return [SeedObjective(value=prompt) for prompt in seed_prompts] +@pytest.fixture +def mock_dataset_config(mock_memory_seeds): + """Create a mock dataset config that returns the seed groups.""" + seed_groups = [SeedAttackGroup(seeds=[seed]) for seed in mock_memory_seeds] + mock_config = MagicMock(spec=DatasetConfiguration) + mock_config.get_all_seed_attack_groups.return_value = seed_groups + mock_config.get_default_dataset_names.return_value = ["airt_leakage"] + mock_config.has_data_source.return_value = True + return mock_config + + @pytest.fixture def first_letter_strategy(): return LeakageStrategy.FIRST_LETTER @@ -211,14 +223,16 @@ class TestLeakageScenarioAttackGeneration: """Tests for LeakageScenario attack generation.""" @pytest.mark.asyncio - async def test_attack_generation_for_all(self, mock_objective_target, mock_objective_scorer, mock_memory_seeds): + async def test_attack_generation_for_all( + self, mock_objective_target, mock_objective_scorer, mock_memory_seeds, mock_dataset_config + ): """Test that _get_atomic_attacks_async returns atomic attacks.""" with patch.object( LeakageScenario, "_get_default_objectives", return_value=[seed.value for seed in mock_memory_seeds] ): scenario = LeakageScenario(objective_scorer=mock_objective_scorer) - await scenario.initialize_async(objective_target=mock_objective_target) + await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) atomic_attacks = await scenario._get_atomic_attacks_async() assert len(atomic_attacks) > 0 @@ -226,7 +240,12 @@ async def test_attack_generation_for_all(self, mock_objective_target, mock_objec @pytest.mark.asyncio async def test_attack_generation_for_first_letter( - self, mock_objective_target, mock_objective_scorer, sample_objectives, first_letter_strategy + self, + mock_objective_target, + mock_objective_scorer, + sample_objectives, + first_letter_strategy, + mock_dataset_config, ): """Test that the first letter attack generation works.""" scenario = LeakageScenario( @@ -235,7 +254,9 @@ async def test_attack_generation_for_first_letter( ) await scenario.initialize_async( - objective_target=mock_objective_target, scenario_strategies=[first_letter_strategy] + objective_target=mock_objective_target, + scenario_strategies=[first_letter_strategy], + dataset_config=mock_dataset_config, ) atomic_attacks = await scenario._get_atomic_attacks_async() for run in atomic_attacks: @@ -243,7 +264,7 @@ async def test_attack_generation_for_first_letter( @pytest.mark.asyncio async def test_attack_generation_for_crescendo( - self, mock_objective_target, mock_objective_scorer, sample_objectives, crescendo_strategy + self, mock_objective_target, mock_objective_scorer, sample_objectives, crescendo_strategy, mock_dataset_config ): """Test that the crescendo attack generation works.""" scenario = LeakageScenario( @@ -252,7 +273,9 @@ async def test_attack_generation_for_crescendo( ) await scenario.initialize_async( - objective_target=mock_objective_target, scenario_strategies=[crescendo_strategy] + objective_target=mock_objective_target, + scenario_strategies=[crescendo_strategy], + dataset_config=mock_dataset_config, ) atomic_attacks = await scenario._get_atomic_attacks_async() @@ -261,7 +284,7 @@ async def test_attack_generation_for_crescendo( @pytest.mark.asyncio async def test_attack_generation_for_image( - self, mock_objective_target, mock_objective_scorer, sample_objectives, image_strategy + self, mock_objective_target, mock_objective_scorer, sample_objectives, image_strategy, mock_dataset_config ): """Test that the image attack generation works.""" scenario = LeakageScenario( @@ -269,14 +292,18 @@ async def test_attack_generation_for_image( objective_scorer=mock_objective_scorer, ) - await scenario.initialize_async(objective_target=mock_objective_target, scenario_strategies=[image_strategy]) + await scenario.initialize_async( + objective_target=mock_objective_target, + scenario_strategies=[image_strategy], + dataset_config=mock_dataset_config, + ) atomic_attacks = await scenario._get_atomic_attacks_async() for run in atomic_attacks: assert isinstance(run._attack, PromptSendingAttack) @pytest.mark.asyncio async def test_attack_generation_for_role_play( - self, mock_objective_target, mock_objective_scorer, sample_objectives, role_play_strategy + self, mock_objective_target, mock_objective_scorer, sample_objectives, role_play_strategy, mock_dataset_config ): """Test that the role play attack generation works.""" scenario = LeakageScenario( @@ -285,7 +312,9 @@ async def test_attack_generation_for_role_play( ) await scenario.initialize_async( - objective_target=mock_objective_target, scenario_strategies=[role_play_strategy] + objective_target=mock_objective_target, + scenario_strategies=[role_play_strategy], + dataset_config=mock_dataset_config, ) atomic_attacks = await scenario._get_atomic_attacks_async() for run in atomic_attacks: @@ -293,7 +322,7 @@ async def test_attack_generation_for_role_play( @pytest.mark.asyncio async def test_attack_runs_include_objectives( - self, mock_objective_target, mock_objective_scorer, sample_objectives + self, mock_objective_target, mock_objective_scorer, sample_objectives, mock_dataset_config ): """Test that attack runs include objectives for each seed prompt.""" scenario = LeakageScenario( @@ -301,7 +330,7 @@ async def test_attack_runs_include_objectives( objective_scorer=mock_objective_scorer, ) - await scenario.initialize_async(objective_target=mock_objective_target) + await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) atomic_attacks = await scenario._get_atomic_attacks_async() # Check that objectives are created for each seed prompt @@ -312,7 +341,7 @@ async def test_attack_runs_include_objectives( @pytest.mark.asyncio async def test_get_atomic_attacks_async_returns_attacks( - self, mock_objective_target, mock_objective_scorer, sample_objectives + self, mock_objective_target, mock_objective_scorer, sample_objectives, mock_dataset_config ): """Test that _get_atomic_attacks_async returns atomic attacks.""" scenario = LeakageScenario( @@ -320,21 +349,21 @@ async def test_get_atomic_attacks_async_returns_attacks( objective_scorer=mock_objective_scorer, ) - await scenario.initialize_async(objective_target=mock_objective_target) + await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) atomic_attacks = await scenario._get_atomic_attacks_async() assert len(atomic_attacks) > 0 assert all(hasattr(run, "_attack") for run in atomic_attacks) @pytest.mark.asyncio async def test_unknown_strategy_raises_value_error( - self, mock_objective_target, mock_objective_scorer, sample_objectives + self, mock_objective_target, mock_objective_scorer, sample_objectives, mock_dataset_config ): """Test that an unknown strategy raises ValueError.""" scenario = LeakageScenario( objectives=sample_objectives, objective_scorer=mock_objective_scorer, ) - await scenario.initialize_async(objective_target=mock_objective_target) + await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) with pytest.raises(ValueError, match="Unknown LeakageStrategy"): await scenario._get_atomic_attack_from_strategy_async("unknown_strategy") @@ -348,19 +377,21 @@ class TestLeakageScenarioLifecycle: @pytest.mark.asyncio async def test_initialize_async_with_max_concurrency( - self, mock_objective_target, mock_objective_scorer, mock_memory_seeds + self, mock_objective_target, mock_objective_scorer, mock_memory_seeds, mock_dataset_config ): """Test initialization with custom max_concurrency.""" with patch.object( LeakageScenario, "_get_default_objectives", return_value=[seed.value for seed in mock_memory_seeds] ): scenario = LeakageScenario(objective_scorer=mock_objective_scorer) - await scenario.initialize_async(objective_target=mock_objective_target, max_concurrency=20) + await scenario.initialize_async( + objective_target=mock_objective_target, max_concurrency=20, dataset_config=mock_dataset_config + ) assert scenario._max_concurrency == 20 @pytest.mark.asyncio async def test_initialize_async_with_memory_labels( - self, mock_objective_target, mock_objective_scorer, mock_memory_seeds + self, mock_objective_target, mock_objective_scorer, mock_memory_seeds, mock_dataset_config ): """Test initialization with memory labels.""" memory_labels = {"test": "leakage", "category": "scenario"} @@ -374,6 +405,7 @@ async def test_initialize_async_with_memory_labels( await scenario.initialize_async( memory_labels=memory_labels, objective_target=mock_objective_target, + dataset_config=mock_dataset_config, ) assert scenario._memory_labels == memory_labels @@ -407,13 +439,13 @@ def test_required_datasets_returns_airt_leakage(self): assert LeakageScenario.required_datasets() == ["airt_leakage"] @pytest.mark.asyncio - async def test_no_target_duplication(self, mock_objective_target, mock_memory_seeds): + async def test_no_target_duplication(self, mock_objective_target, mock_memory_seeds, mock_dataset_config): """Test that all three targets (adversarial, object, scorer) are distinct.""" with patch.object( LeakageScenario, "_get_default_objectives", return_value=[seed.value for seed in mock_memory_seeds] ): scenario = LeakageScenario() - await scenario.initialize_async(objective_target=mock_objective_target) + await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) objective_target = scenario._objective_target @@ -565,7 +597,7 @@ def test_ensure_blank_image_exists_creates_parent_directories( @pytest.mark.asyncio async def test_image_strategy_uses_add_image_text_converter( - self, mock_objective_target, mock_objective_scorer, sample_objectives, image_strategy + self, mock_objective_target, mock_objective_scorer, sample_objectives, image_strategy, mock_dataset_config ): """Test that the image strategy uses AddImageTextConverter (not AddTextImageConverter).""" from pyrit.prompt_converter import AddImageTextConverter @@ -575,7 +607,11 @@ async def test_image_strategy_uses_add_image_text_converter( objective_scorer=mock_objective_scorer, ) - await scenario.initialize_async(objective_target=mock_objective_target, scenario_strategies=[image_strategy]) + await scenario.initialize_async( + objective_target=mock_objective_target, + scenario_strategies=[image_strategy], + dataset_config=mock_dataset_config, + ) atomic_attacks = await scenario._get_atomic_attacks_async() # Verify the attack uses AddImageTextConverter diff --git a/tests/unit/scenarios/test_scam.py b/tests/unit/scenarios/test_scam.py index 8d0aa2f93..f8fcd7975 100644 --- a/tests/unit/scenarios/test_scam.py +++ b/tests/unit/scenarios/test_scam.py @@ -17,8 +17,9 @@ ) from pyrit.executor.attack.core.attack_config import AttackScoringConfig from pyrit.identifiers import ScorerIdentifier -from pyrit.models import SeedDataset, SeedGroup, SeedObjective +from pyrit.models import SeedAttackGroup, SeedDataset, SeedGroup, SeedObjective from pyrit.prompt_target import OpenAIChatTarget, PromptChatTarget, PromptTarget +from pyrit.scenario import DatasetConfiguration from pyrit.scenario.scenarios.airt.scam import Scam, ScamStrategy from pyrit.score import TrueFalseCompositeScorer @@ -42,6 +43,23 @@ def mock_memory_seed_groups() -> List[SeedGroup]: return [SeedGroup(seeds=[SeedObjective(value=prompt)]) for prompt in SEED_PROMPT_LIST] +@pytest.fixture +def mock_memory_seeds(): + """Create mock seeds (SeedObjective objects) from the seed prompt list.""" + return [SeedObjective(value=prompt) for prompt in SEED_PROMPT_LIST] + + +@pytest.fixture +def mock_dataset_config(mock_memory_seed_groups): + """Create a mock dataset config that returns the seed groups.""" + seed_attack_groups = [SeedAttackGroup(seeds=list(sg.seeds)) for sg in mock_memory_seed_groups] + mock_config = MagicMock(spec=DatasetConfiguration) + mock_config.get_all_seed_attack_groups.return_value = seed_attack_groups + mock_config.get_default_dataset_names.return_value = ["airt_scam"] + mock_config.has_data_source.return_value = True + return mock_config + + @pytest.fixture def single_turn_strategy() -> ScamStrategy: return ScamStrategy.SINGLE_TURN @@ -192,13 +210,13 @@ class TestScamAttackGeneration: @pytest.mark.asyncio async def test_attack_generation_for_all( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups, mock_dataset_config ): """Test that _get_atomic_attacks_async returns atomic attacks.""" with patch.object(Scam, "_resolve_seed_groups", return_value=mock_memory_seed_groups): scenario = Scam(objective_scorer=mock_objective_scorer) - await scenario.initialize_async(objective_target=mock_objective_target) + await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) atomic_attacks = await scenario._get_atomic_attacks_async() assert len(atomic_attacks) > 0 @@ -211,16 +229,17 @@ async def test_attack_generation_for_singleturn_async( mock_objective_target: PromptTarget, mock_objective_scorer: TrueFalseCompositeScorer, single_turn_strategy: ScamStrategy, - sample_objectives: List[str], + mock_dataset_config: DatasetConfiguration, ) -> None: """Test that the single turn strategy attack generation works.""" scenario = Scam( - objectives=sample_objectives, objective_scorer=mock_objective_scorer, ) await scenario.initialize_async( - objective_target=mock_objective_target, scenario_strategies=[single_turn_strategy] + objective_target=mock_objective_target, + scenario_strategies=[single_turn_strategy], + dataset_config=mock_dataset_config, ) atomic_attacks = await scenario._get_atomic_attacks_async() @@ -229,16 +248,17 @@ async def test_attack_generation_for_singleturn_async( @pytest.mark.asyncio async def test_attack_generation_for_multiturn_async( - self, mock_objective_target, mock_objective_scorer, sample_objectives, multi_turn_strategy + self, mock_objective_target, mock_objective_scorer, multi_turn_strategy, mock_dataset_config ): """Test that the multi turn attack generation works.""" scenario = Scam( - objectives=sample_objectives, objective_scorer=mock_objective_scorer, ) await scenario.initialize_async( - objective_target=mock_objective_target, scenario_strategies=[multi_turn_strategy] + objective_target=mock_objective_target, + scenario_strategies=[multi_turn_strategy], + dataset_config=mock_dataset_config, ) atomic_attacks = await scenario._get_atomic_attacks_async() @@ -251,21 +271,21 @@ async def test_attack_runs_include_objectives_async( *, mock_objective_target: PromptTarget, mock_objective_scorer: TrueFalseCompositeScorer, - sample_objectives: List[str], + mock_dataset_config: DatasetConfiguration, + mock_memory_seeds, ) -> None: """Test that attack runs include objectives for each seed prompt.""" scenario = Scam( - objectives=sample_objectives, objective_scorer=mock_objective_scorer, ) - await scenario.initialize_async(objective_target=mock_objective_target) + await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) atomic_attacks = await scenario._get_atomic_attacks_async() for run in atomic_attacks: - assert len(run.objectives) == len(sample_objectives) + assert len(run.objectives) == len(mock_memory_seeds) for index, objective in enumerate(run.objectives): - assert sample_objectives[index] in objective + assert mock_memory_seeds[index].value in objective @pytest.mark.asyncio async def test_get_atomic_attacks_async_returns_attacks( @@ -273,15 +293,14 @@ async def test_get_atomic_attacks_async_returns_attacks( *, mock_objective_target: PromptTarget, mock_objective_scorer: TrueFalseCompositeScorer, - sample_objectives: List[str], + mock_dataset_config: DatasetConfiguration, ) -> None: """Test that _get_atomic_attacks_async returns atomic attacks.""" scenario = Scam( - objectives=sample_objectives, objective_scorer=mock_objective_scorer, ) - await scenario.initialize_async(objective_target=mock_objective_target) + await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) atomic_attacks = await scenario._get_atomic_attacks_async() assert len(atomic_attacks) > 0 assert all(hasattr(run, "_attack") for run in atomic_attacks) @@ -298,11 +317,14 @@ async def test_initialize_async_with_max_concurrency( mock_objective_target: PromptTarget, mock_objective_scorer: TrueFalseCompositeScorer, mock_memory_seed_groups: List[SeedGroup], + mock_dataset_config, ) -> None: """Test initialization with custom max_concurrency.""" with patch.object(Scam, "_resolve_seed_groups", return_value=mock_memory_seed_groups): scenario = Scam(objective_scorer=mock_objective_scorer) - await scenario.initialize_async(objective_target=mock_objective_target, max_concurrency=20) + await scenario.initialize_async( + objective_target=mock_objective_target, max_concurrency=20, dataset_config=mock_dataset_config + ) assert scenario._max_concurrency == 20 @pytest.mark.asyncio @@ -312,6 +334,7 @@ async def test_initialize_async_with_memory_labels( mock_objective_target: PromptTarget, mock_objective_scorer: TrueFalseCompositeScorer, mock_memory_seed_groups: List[SeedGroup], + mock_dataset_config, ) -> None: """Test initialization with memory labels.""" memory_labels = {"type": "scam", "category": "scenario"} @@ -321,6 +344,7 @@ async def test_initialize_async_with_memory_labels( await scenario.initialize_async( memory_labels=memory_labels, objective_target=mock_objective_target, + dataset_config=mock_dataset_config, ) assert scenario._memory_labels == memory_labels @@ -333,11 +357,9 @@ def test_scenario_version_is_set( self, *, mock_objective_scorer: TrueFalseCompositeScorer, - sample_objectives: List[str], ) -> None: """Test that scenario version is properly set.""" scenario = Scam( - objectives=sample_objectives, objective_scorer=mock_objective_scorer, ) @@ -345,12 +367,12 @@ def test_scenario_version_is_set( @pytest.mark.asyncio async def test_no_target_duplication_async( - self, *, mock_objective_target: PromptTarget, mock_memory_seed_groups: List[SeedGroup] + self, *, mock_objective_target: PromptTarget, mock_memory_seed_groups: List[SeedGroup], mock_dataset_config ) -> None: """Test that all three targets (adversarial, object, scorer) are distinct.""" with patch.object(Scam, "_resolve_seed_groups", return_value=mock_memory_seed_groups): scenario = Scam() - await scenario.initialize_async(objective_target=mock_objective_target) + await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) objective_target = scenario._objective_target scorer_target = scenario._scorer_config.objective_scorer # type: ignore diff --git a/tests/unit/scenarios/test_scenario.py b/tests/unit/scenarios/test_scenario.py index 909038cdf..9b13d8730 100644 --- a/tests/unit/scenarios/test_scenario.py +++ b/tests/unit/scenarios/test_scenario.py @@ -71,7 +71,10 @@ def mock_atomic_attacks(): def mock_objective_target(): """Create a mock objective target for testing.""" target = MagicMock() - target.get_identifier.return_value = {"__type__": "MockTarget", "__module__": "test"} + target.get_identifier.return_value = { + "__type__": "MockTarget", + "__module__": "test", + } return target @@ -82,7 +85,11 @@ def sample_attack_results(): AttackResult( conversation_id=f"conv-{i}", objective=f"objective{i}", - attack_identifier={"__type__": "TestAttack", "__module__": "test", "id": str(i)}, + attack_identifier={ + "__type__": "TestAttack", + "__module__": "test", + "id": str(i), + }, outcome=AttackOutcome.SUCCESS, executed_turns=1, ) @@ -223,7 +230,10 @@ async def test_initialize_async_sets_objective_target(self, mock_objective_targe await scenario.initialize_async(objective_target=mock_objective_target) assert scenario._objective_target == mock_objective_target - assert scenario._objective_target_identifier == {"__type__": "MockTarget", "__module__": "test"} + assert scenario._objective_target_identifier == { + "__type__": "MockTarget", + "__module__": "test", + } @pytest.mark.asyncio async def test_initialize_async_requires_objective_target(self): @@ -432,7 +442,11 @@ async def test_run_async_returns_scenario_result_with_identifier( assert result.scenario_identifier.name == "ConcreteScenario" assert result.scenario_identifier.version == 5 assert result.scenario_identifier.pyrit_version is not None - assert result.get_strategies_used() == ["attack_run_1", "attack_run_2", "attack_run_3"] + assert result.get_strategies_used() == [ + "attack_run_1", + "attack_run_2", + "attack_run_3", + ] @pytest.mark.usefixtures("patch_central_database") @@ -527,7 +541,10 @@ def test_scenario_result_with_empty_results(self): identifier = ScenarioIdentifier(name="TestScenario", scenario_version=1) result = ScenarioResult( scenario_identifier=identifier, - objective_target_identifier={"__type__": "TestTarget", "__module__": "test"}, + objective_target_identifier={ + "__type__": "TestTarget", + "__module__": "test", + }, attack_results={"base64": []}, objective_scorer_identifier=_TEST_SCORER_ID, ) @@ -542,7 +559,10 @@ def test_scenario_result_objective_achieved_rate(self, sample_attack_results): # All successful result = ScenarioResult( scenario_identifier=identifier, - objective_target_identifier={"__type__": "TestTarget", "__module__": "test"}, + objective_target_identifier={ + "__type__": "TestTarget", + "__module__": "test", + }, attack_results={"base64": sample_attack_results}, objective_scorer_identifier=_TEST_SCORER_ID, ) @@ -553,21 +573,32 @@ def test_scenario_result_objective_achieved_rate(self, sample_attack_results): AttackResult( conversation_id="conv-fail", objective="objective", - attack_identifier={"__type__": "TestAttack", "__module__": "test", "id": "1"}, + attack_identifier={ + "__type__": "TestAttack", + "__module__": "test", + "id": "1", + }, outcome=AttackOutcome.FAILURE, executed_turns=1, ), AttackResult( conversation_id="conv-fail2", objective="objective", - attack_identifier={"__type__": "TestAttack", "__module__": "test", "id": "2"}, + attack_identifier={ + "__type__": "TestAttack", + "__module__": "test", + "id": "2", + }, outcome=AttackOutcome.FAILURE, executed_turns=1, ), ] result2 = ScenarioResult( scenario_identifier=identifier, - objective_target_identifier={"__type__": "TestTarget", "__module__": "test"}, + objective_target_identifier={ + "__type__": "TestTarget", + "__module__": "test", + }, attack_results={"base64": mixed_results}, objective_scorer_identifier=_TEST_SCORER_ID, ) @@ -599,3 +630,197 @@ def test_scenario_identifier_with_init_data(self): identifier = ScenarioIdentifier(name="TestScenario", scenario_version=1, init_data=init_data) assert identifier.init_data == init_data + + +def create_mock_truefalse_scorer(): + """Create a mock TrueFalseScorer for testing baseline-only execution.""" + from pyrit.score import TrueFalseScorer + + mock_scorer = MagicMock(spec=TrueFalseScorer) + mock_scorer.get_identifier.return_value = { + "__type__": "MockTrueFalseScorer", + "__module__": "test", + } + mock_scorer.get_scorer_metrics.return_value = None + # Make isinstance check work + mock_scorer.__class__ = TrueFalseScorer + return mock_scorer + + +class ConcreteScenarioWithTrueFalseScorer(Scenario): + """Concrete implementation of Scenario for testing baseline-only execution.""" + + def __init__(self, atomic_attacks_to_return=None, **kwargs): + # Add required strategy_class if not provided + + class TestStrategy(ScenarioStrategy): + TEST = ("test", {"concrete"}) + ALL = ("all", {"all"}) + + @classmethod + def get_aggregate_tags(cls) -> set[str]: + return {"all"} + + kwargs.setdefault("strategy_class", TestStrategy) + + # Use TrueFalseScorer mock if not provided + if "objective_scorer" not in kwargs: + kwargs["objective_scorer"] = create_mock_truefalse_scorer() + + super().__init__(**kwargs) + self._atomic_attacks_to_return = atomic_attacks_to_return or [] + + @classmethod + def get_strategy_class(cls): + """Return a mock strategy class for testing.""" + + from pyrit.scenario.core.scenario_strategy import ScenarioStrategy + + class TestStrategy(ScenarioStrategy): + TEST = ("test", {"concrete"}) + ALL = ("all", {"all"}) + + @classmethod + def get_aggregate_tags(cls) -> set[str]: + return {"all"} + + return TestStrategy + + @classmethod + def get_default_strategy(cls): + """Return the default strategy for testing.""" + return cls.get_strategy_class().ALL + + @classmethod + def default_dataset_config(cls) -> DatasetConfiguration: + """Return the default dataset configuration for testing.""" + return DatasetConfiguration() + + async def _get_atomic_attacks_async(self): + return self._atomic_attacks_to_return + + +@pytest.mark.usefixtures("patch_central_database") +class TestScenarioBaselineOnlyExecution: + """Tests for baseline-only execution (empty strategies with include_baseline=True).""" + + @pytest.mark.asyncio + async def test_initialize_async_with_empty_strategies_and_baseline(self, mock_objective_target): + """Test that baseline-only execution works when include_baseline=True and strategies is empty.""" + from pyrit.models import SeedAttackGroup, SeedObjective + + # Create a scenario with include_default_baseline=True and TrueFalseScorer + scenario = ConcreteScenarioWithTrueFalseScorer( + name="Baseline Only Test", + version=1, + include_default_baseline=True, + ) + + # Create a mock dataset config with seed groups + mock_dataset_config = MagicMock(spec=DatasetConfiguration) + mock_dataset_config.get_all_seed_attack_groups.return_value = [ + SeedAttackGroup(seeds=[SeedObjective(value="test objective 1")]), + SeedAttackGroup(seeds=[SeedObjective(value="test objective 2")]), + ] + + # Initialize with empty strategies + await scenario.initialize_async( + objective_target=mock_objective_target, + scenario_strategies=[], # Empty list - baseline only + dataset_config=mock_dataset_config, + ) + + # Should have exactly one attack - the baseline + assert scenario.atomic_attack_count == 1 + assert scenario._atomic_attacks[0].atomic_attack_name == "baseline" + + @pytest.mark.asyncio + async def test_baseline_only_execution_runs_successfully(self, mock_objective_target, sample_attack_results): + """Test that baseline-only scenario can run successfully.""" + from pyrit.models import SeedAttackGroup, SeedObjective + + # Create a scenario with include_default_baseline=True and TrueFalseScorer + scenario = ConcreteScenarioWithTrueFalseScorer( + name="Baseline Only Test", + version=1, + include_default_baseline=True, + ) + + # Create a mock dataset config with seed groups + mock_dataset_config = MagicMock(spec=DatasetConfiguration) + mock_dataset_config.get_all_seed_attack_groups.return_value = [ + SeedAttackGroup(seeds=[SeedObjective(value="test objective 1")]), + ] + + # Initialize with empty strategies + await scenario.initialize_async( + objective_target=mock_objective_target, + scenario_strategies=[], # Empty list - baseline only + dataset_config=mock_dataset_config, + ) + + # Mock the baseline attack's run_async + scenario._atomic_attacks[0].run_async = create_mock_run_async([sample_attack_results[0]]) + + # Run the scenario + result = await scenario.run_async() + + # Verify the result + assert isinstance(result, ScenarioResult) + assert "baseline" in result.attack_results + assert len(result.attack_results["baseline"]) == 1 + + @pytest.mark.asyncio + async def test_empty_strategies_without_baseline_allows_initialization(self, mock_objective_target): + """Test that empty strategies without include_baseline allows initialization but fails at run time.""" + scenario = ConcreteScenario( + name="No Baseline Test", + version=1, + include_default_baseline=False, # No baseline + ) + + mock_dataset_config = MagicMock(spec=DatasetConfiguration) + + # Empty strategies are now always allowed during initialization + # (no allow_empty parameter required) + await scenario.initialize_async( + objective_target=mock_objective_target, + scenario_strategies=[], # Empty list without baseline + dataset_config=mock_dataset_config, + ) + + # But running should fail because there are no atomic attacks + with pytest.raises(ValueError, match="Cannot run scenario with no atomic attacks"): + await scenario.run_async() + + @pytest.mark.asyncio + async def test_standalone_baseline_uses_dataset_config_seeds(self, mock_objective_target): + """Test that standalone baseline uses seed groups from dataset_config.""" + from pyrit.models import SeedAttackGroup, SeedObjective + + scenario = ConcreteScenarioWithTrueFalseScorer( + name="Baseline Seeds Test", + version=1, + include_default_baseline=True, + ) + + # Create specific seed groups to verify they're used + expected_seeds = [ + SeedAttackGroup(seeds=[SeedObjective(value="objective_a")]), + SeedAttackGroup(seeds=[SeedObjective(value="objective_b")]), + SeedAttackGroup(seeds=[SeedObjective(value="objective_c")]), + ] + + mock_dataset_config = MagicMock(spec=DatasetConfiguration) + mock_dataset_config.get_all_seed_attack_groups.return_value = expected_seeds + + await scenario.initialize_async( + objective_target=mock_objective_target, + scenario_strategies=[], + dataset_config=mock_dataset_config, + ) + + # Verify the baseline attack has the expected seed groups + baseline_attack = scenario._atomic_attacks[0] + assert baseline_attack.atomic_attack_name == "baseline" + assert baseline_attack.seed_groups == expected_seeds