# Copyright (c) Microsoft Corporation. # Licensed under the MIT license. """Tests for the scenarios.AtomicAttack class.""" import inspect from unittest.mock import AsyncMock, MagicMock, patch import pytest from pyrit.executor.attack import AttackExecutor, AttackStrategy from pyrit.executor.attack.core import AttackExecutorResult from pyrit.models import ( AtomicAttackIdentifier, AttackOutcome, AttackResult, AttackSeedGroup, ComponentIdentifier, SeedGroup, SeedObjective, SeedPrompt, ) from pyrit.scenario import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique @pytest.fixture def mock_attack(): """Create a mock AttackStrategy for testing.""" attack = MagicMock(spec=AttackStrategy) attack.get_identifier.return_value = ComponentIdentifier(class_name="MockAttack", class_module="pyrit.test") return attack @pytest.fixture def sample_seed_groups(): """Create sample seed groups with objectives for testing.""" return [ AttackSeedGroup( seeds=[ SeedObjective(value="objective1"), SeedPrompt(value="prompt1"), ] ), AttackSeedGroup( seeds=[ SeedObjective(value="objective2"), SeedPrompt(value="prompt2"), ] ), AttackSeedGroup( seeds=[ SeedObjective(value="objective3"), SeedPrompt(value="prompt3"), ] ), ] @pytest.fixture def sample_seed_groups_without_objectives(): """Create sample seed groups without objectives for testing. Note: AttackSeedGroup now validates exactly one objective at construction, so we use SeedGroup here which doesn't have that requirement. """ return [ SeedGroup( seeds=[ SeedPrompt(value="prompt1"), ] ), ] @pytest.fixture def sample_attack_results(): """Create sample attack results for testing.""" return [ AttackResult( conversation_id="conv-1", objective="objective1", outcome=AttackOutcome.SUCCESS, executed_turns=1, ), AttackResult( conversation_id="conv-2", objective="objective2", outcome=AttackOutcome.SUCCESS, executed_turns=1, ), AttackResult( conversation_id="conv-3", objective="objective3", outcome=AttackOutcome.FAILURE, executed_turns=1, ), ] def wrap_results(results): """Helper to wrap attack results in AttackExecutorResult.""" return AttackExecutorResult( completed_results=results, incomplete_objectives=[], input_indices=list(range(len(results))), ) @pytest.mark.usefixtures("patch_central_database") class TestAtomicAttackInitialization: """Tests for AtomicAttack class initialization.""" def test_init_with_valid_params(self, mock_attack, sample_seed_groups): """Test successful initialization with valid parameters.""" atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="Test Attack Run", ) assert atomic_attack._attack_technique.attack == mock_attack assert atomic_attack._seed_groups == sample_seed_groups assert atomic_attack._memory_labels == {} assert atomic_attack._attack_execute_params == {} def test_init_with_memory_labels(self, mock_attack, sample_seed_groups): """Test initialization with memory labels.""" memory_labels = {"test": "label", "category": "attack"} atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, memory_labels=memory_labels, atomic_attack_name="Test Attack Run", ) assert atomic_attack._memory_labels == memory_labels def test_init_with_attack_execute_params(self, mock_attack, sample_seed_groups): """Test initialization with additional attack execute parameters.""" atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, max_retries=5, custom_param="value", atomic_attack_name="Test Attack Run", ) assert atomic_attack._attack_execute_params["max_retries"] == 5 assert atomic_attack._attack_execute_params["custom_param"] == "value" def test_init_with_all_parameters(self, mock_attack, sample_seed_groups): """Test initialization with all parameters.""" memory_labels = {"test": "comprehensive"} atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, memory_labels=memory_labels, batch_size=10, timeout=30, atomic_attack_name="Test Attack Run", ) assert atomic_attack._attack_technique.attack == mock_attack assert atomic_attack._seed_groups == sample_seed_groups assert atomic_attack._memory_labels == memory_labels assert atomic_attack._attack_execute_params["batch_size"] == 10 assert atomic_attack._attack_execute_params["timeout"] == 30 def test_init_fails_with_empty_seed_groups(self, mock_attack): """Test that initialization fails when seed_groups list is empty.""" with pytest.raises(ValueError, match="seed_groups list cannot be empty"): AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=[], atomic_attack_name="Test Attack Run", ) def test_init_fails_with_seed_group_missing_objective(self, mock_attack): """Test that AttackSeedGroup without objective cannot be created. AttackSeedGroup now validates exactly one objective at construction time, so we can't even create one without an objective. """ # AttackSeedGroup now validates exactly one objective at construction with pytest.raises(ValueError, match="must have exactly one objective"): AttackSeedGroup(seeds=[SeedPrompt(value="prompt1")]) def test_objectives_property_returns_values_from_seed_groups(self, mock_attack, sample_seed_groups): """Test that the objectives property returns values from seed groups.""" atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="Test Attack Run", ) assert atomic_attack.objectives == ["objective1", "objective2", "objective3"] def test_seed_groups_property_returns_copy(self, mock_attack, sample_seed_groups): """Test that the seed_groups property returns a copy.""" atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="Test Attack Run", ) returned_groups = atomic_attack.seed_groups assert returned_groups == sample_seed_groups assert returned_groups is not atomic_attack._seed_groups @pytest.mark.usefixtures("patch_central_database") class TestAtomicAttackExecution: """Tests for AtomicAttack execution methods.""" async def test_run_async_with_valid_atomic_attack(self, mock_attack, sample_seed_groups, sample_attack_results): """Test successful execution of an atomic attack.""" atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="Test Attack Run", ) with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results(sample_attack_results) result = await atomic_attack.run_async() assert len(result.completed_results) == 3 assert result.completed_results == sample_attack_results assert len(result.incomplete_objectives) == 0 mock_exec.assert_called_once() # Verify the attack was passed correctly call_kwargs = mock_exec.call_args.kwargs assert call_kwargs["attack"] == mock_attack async def test_run_async_with_default_concurrency(self, mock_attack, sample_seed_groups, sample_attack_results): """Test that default concurrency (1) is used when not specified.""" atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="Test Attack Run", ) with ( patch.object(AttackExecutor, "__init__", return_value=None) as mock_init, patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec, ): mock_exec.return_value = wrap_results(sample_attack_results) await atomic_attack.run_async() mock_init.assert_called_once_with(max_concurrency=1) async def test_run_async_with_injected_executor_reuses_it( self, mock_attack, sample_seed_groups, sample_attack_results ): """When an executor is passed, AtomicAttack must reuse it rather than build a new one.""" atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="Test Attack Run", ) injected = AttackExecutor(max_concurrency=7) with ( patch.object(AttackExecutor, "__init__", return_value=None) as mock_init, patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec, ): mock_exec.return_value = wrap_results(sample_attack_results) await atomic_attack.run_async(executor=injected) # __init__ must not be called again — the injected executor is reused as-is. mock_init.assert_not_called() async def test_run_async_passes_memory_labels(self, mock_attack, sample_seed_groups, sample_attack_results): """Test that memory labels are passed to the executor.""" memory_labels = {"test": "attack_run", "category": "attack"} atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, memory_labels=memory_labels, atomic_attack_name="Test Attack Run", ) with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results(sample_attack_results) await atomic_attack.run_async() call_kwargs = mock_exec.call_args.kwargs assert "memory_labels" in call_kwargs assert call_kwargs["memory_labels"] == memory_labels async def test_run_async_passes_seed_groups(self, mock_attack, sample_seed_groups, sample_attack_results): """Test that seed_groups are passed to the executor.""" atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="Test Attack Run", ) with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results(sample_attack_results) await atomic_attack.run_async() call_kwargs = mock_exec.call_args.kwargs assert "seed_groups" in call_kwargs assert call_kwargs["seed_groups"] == sample_seed_groups async def test_run_async_raises_when_executor_returns_non_attack_result(self, mock_attack, sample_seed_groups): """Test that a non-AttackResult item from the executor raises ValueError.""" atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="Test Attack Run", ) with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results(["not-an-attack-result"]) with pytest.raises(ValueError, match="unsupported result type"): await atomic_attack.run_async() async def test_run_async_passes_attack_execute_params(self, mock_attack, sample_seed_groups, sample_attack_results): """Test that attack execute parameters are passed to the executor.""" atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, custom_param="value", max_retries=3, atomic_attack_name="Test Attack Run", ) with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results(sample_attack_results) await atomic_attack.run_async() call_kwargs = mock_exec.call_args.kwargs assert call_kwargs["custom_param"] == "value" assert call_kwargs["max_retries"] == 3 async def test_run_async_merges_all_parameters(self, mock_attack, sample_seed_groups, sample_attack_results): """Test that all parameters are merged and passed correctly.""" memory_labels = {"test": "merge"} atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, memory_labels=memory_labels, batch_size=5, atomic_attack_name="Test Attack Run", ) with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results(sample_attack_results) await atomic_attack.run_async() call_kwargs = mock_exec.call_args.kwargs assert call_kwargs["attack"] == mock_attack assert call_kwargs["seed_groups"] == sample_seed_groups assert call_kwargs["memory_labels"] == memory_labels assert call_kwargs["batch_size"] == 5 async def test_run_async_handles_execution_failure(self, mock_attack, sample_seed_groups): """Test that execution failures are properly handled and raised.""" atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="Test Attack Run", ) with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.side_effect = Exception("Execution error") with pytest.raises(ValueError, match="Failed to execute atomic attack 'Test Attack Run'"): await atomic_attack.run_async() async def test_run_async_passes_return_partial_on_failure_true_by_default( self, mock_attack, sample_seed_groups, sample_attack_results ): """Test that atomic attack passes return_partial_on_failure=True by default.""" atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="Test Attack Run", ) with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results(sample_attack_results) await atomic_attack.run_async() call_kwargs = mock_exec.call_args.kwargs assert "return_partial_on_failure" in call_kwargs assert call_kwargs["return_partial_on_failure"] is True async def test_run_async_respects_explicit_return_partial_on_failure( self, mock_attack, sample_seed_groups, sample_attack_results ): """Test that explicit return_partial_on_failure parameter is passed through.""" atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="Test Attack Run", ) with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results(sample_attack_results) await atomic_attack.run_async(return_partial_on_failure=False) call_kwargs = mock_exec.call_args.kwargs assert "return_partial_on_failure" in call_kwargs assert call_kwargs["return_partial_on_failure"] is False @pytest.mark.usefixtures("patch_central_database") class TestAtomicAttackIntegration: """Integration Tests for AtomicAttack.""" async def test_full_attack_run_execution_flow(self, mock_attack, sample_seed_groups): """Test the complete attack run execution flow end-to-end.""" memory_labels = {"test": "integration", "attack_run": "full"} atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, memory_labels=memory_labels, batch_size=2, atomic_attack_name="Test Attack Run", ) mock_results = [ AttackResult( conversation_id=f"conv-{i}", objective=f"objective{i + 1}", outcome=AttackOutcome.SUCCESS, executed_turns=1, ) for i in range(3) ] with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results(mock_results) attack_run_result = await atomic_attack.run_async() assert len(attack_run_result.completed_results) == 3 for i, result in enumerate(attack_run_result.completed_results): assert result.objective == f"objective{i + 1}" assert result.outcome == AttackOutcome.SUCCESS call_kwargs = mock_exec.call_args.kwargs assert call_kwargs["attack"] == mock_attack assert call_kwargs["seed_groups"] == sample_seed_groups assert call_kwargs["memory_labels"] == memory_labels assert call_kwargs["batch_size"] == 2 async def test_atomic_attack_with_single_seed_group(self, mock_attack): """Test atomic attack with a single seed group.""" single_seed_group = [ AttackSeedGroup( seeds=[ SeedObjective(value="single_objective"), SeedPrompt(value="single_prompt"), ] ) ] atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=single_seed_group, atomic_attack_name="Test Attack Run", ) mock_result = [ AttackResult( conversation_id="conv-1", objective="single_objective", outcome=AttackOutcome.SUCCESS, executed_turns=1, ) ] with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results(mock_result) attack_run_result = await atomic_attack.run_async() assert len(attack_run_result.completed_results) == 1 assert attack_run_result.completed_results[0].objective == "single_objective" async def test_atomic_attack_with_many_seed_groups(self, mock_attack): """Test atomic attack with many seed groups.""" many_seed_groups = [ AttackSeedGroup( seeds=[ SeedObjective(value=f"objective_{i}"), SeedPrompt(value=f"prompt_{i}"), ] ) for i in range(20) ] atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=many_seed_groups, atomic_attack_name="Test Attack Run", ) mock_results = [ AttackResult( conversation_id=f"conv-{i}", objective=f"objective_{i}", outcome=AttackOutcome.SUCCESS, executed_turns=1, ) for i in range(20) ] with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results(mock_results) attack_run_result = await atomic_attack.run_async() assert len(attack_run_result.completed_results) == 20 call_kwargs = mock_exec.call_args.kwargs assert len(call_kwargs["seed_groups"]) == 20 @pytest.mark.usefixtures("patch_central_database") class TestAtomicAttackExecutorParamCompatibility: """Tests to verify AtomicAttack passes parameters compatible with AttackExecutor.""" def test_atomic_attack_passes_expected_executor_params(self, mock_attack, sample_seed_groups): """ Test that AtomicAttack.run_async passes all expected parameters to execute_attack_from_seed_groups_async. """ # Get the signature of execute_attack_from_seed_groups_async executor_method = AttackExecutor.execute_attack_from_seed_groups_async sig = inspect.signature(executor_method) # These are the parameters that execute_attack_from_seed_groups_async accepts expected_params = set(sig.parameters.keys()) - {"self"} # Verify the explicit parameters we know AtomicAttack should pass # Note: memory_labels is passed via **broadcast_fields, not as an explicit parameter required_from_atomic_attack = { "attack", "seed_groups", "return_partial_on_failure", } # All required params should be in the executor method signature assert required_from_atomic_attack.issubset(expected_params), ( f"Missing expected params in executor: {required_from_atomic_attack - expected_params}" ) # Verify that the executor accepts **broadcast_fields (e.g., for memory_labels) assert "broadcast_fields" in expected_params, "Executor should accept **broadcast_fields for dynamic params" async def test_run_async_only_passes_valid_executor_params( self, mock_attack, sample_seed_groups, sample_attack_results ): """ Test that run_async doesn't pass parameters that the executor doesn't accept. The executor has strict_param_matching so invalid params would cause failures. """ atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="Test Attack Run", ) with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results(sample_attack_results) await atomic_attack.run_async() call_kwargs = mock_exec.call_args.kwargs # Verify essential params are present assert "attack" in call_kwargs assert "seed_groups" in call_kwargs assert "memory_labels" in call_kwargs assert "return_partial_on_failure" in call_kwargs @pytest.mark.usefixtures("patch_central_database") class TestAtomicAttackWithMessages: """Tests for AtomicAttack with seed groups containing multi-turn messages.""" @pytest.fixture def seed_groups_with_messages(self): """Create seed groups with multi-turn message sequences for testing.""" return [ AttackSeedGroup( seeds=[ SeedObjective(value="multi_turn_objective_1"), SeedPrompt(value="First message", data_type="text", sequence=0, role="user"), SeedPrompt(value="Second message", data_type="text", sequence=1, role="user"), SeedPrompt(value="Third message", data_type="text", sequence=2, role="user"), ] ), AttackSeedGroup( seeds=[ SeedObjective(value="multi_turn_objective_2"), SeedPrompt(value="Message A", data_type="text", sequence=0, role="user"), SeedPrompt(value="Message B", data_type="text", sequence=1, role="user"), ] ), ] @pytest.fixture def mixed_seed_groups(self): """Create seed groups where some have messages and some don't.""" return [ # No messages (just objective) AttackSeedGroup(seeds=[SeedObjective(value="simple_objective")]), # With messages - roles required for multi-sequence AttackSeedGroup( seeds=[ SeedObjective(value="objective_with_messages"), SeedPrompt(value="Message 1", data_type="text", sequence=0, role="user"), SeedPrompt(value="Message 2", data_type="text", sequence=1, role="user"), ] ), ] def test_init_with_seed_groups_with_messages(self, mock_attack, seed_groups_with_messages): """Test that AtomicAttack initializes correctly with seed groups containing messages.""" atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=seed_groups_with_messages, atomic_attack_name="Multi-turn Attack", ) assert len(atomic_attack.seed_groups) == 2 assert atomic_attack.objectives == ["multi_turn_objective_1", "multi_turn_objective_2"] # Verify seed groups have user messages for sg in atomic_attack.seed_groups: assert len(sg.user_messages) > 0 def test_seed_groups_user_messages_property(self, mock_attack, seed_groups_with_messages): """Test that seed group user_messages are accessible and have correct content.""" atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=seed_groups_with_messages, atomic_attack_name="Multi-turn Attack", ) sg1 = atomic_attack.seed_groups[0] sg2 = atomic_attack.seed_groups[1] # First seed group has 3 user messages assert len(sg1.user_messages) == 3 assert sg1.user_messages[0].message_pieces[0].original_value == "First message" assert sg1.user_messages[1].message_pieces[0].original_value == "Second message" assert sg1.user_messages[2].message_pieces[0].original_value == "Third message" # Second seed group has 2 user messages assert len(sg2.user_messages) == 2 assert sg2.user_messages[0].message_pieces[0].original_value == "Message A" assert sg2.user_messages[1].message_pieces[0].original_value == "Message B" async def test_run_async_passes_seed_groups_with_messages(self, mock_attack, seed_groups_with_messages): """Test that run_async correctly passes seed groups with messages to executor.""" atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=seed_groups_with_messages, atomic_attack_name="Multi-turn Attack", ) mock_results = [ AttackResult( conversation_id=f"conv-{i}", objective=seed_groups_with_messages[i].objective.value, outcome=AttackOutcome.SUCCESS, executed_turns=len(seed_groups_with_messages[i].user_messages), ) for i in range(2) ] with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results(mock_results) result = await atomic_attack.run_async() assert len(result.completed_results) == 2 # Verify seed groups were passed correctly call_kwargs = mock_exec.call_args.kwargs passed_seed_groups = call_kwargs["seed_groups"] assert len(passed_seed_groups) == 2 # Verify user messages are preserved in passed seed groups assert len(passed_seed_groups[0].user_messages) == 3 assert len(passed_seed_groups[1].user_messages) == 2 def test_init_with_mixed_seed_groups(self, mock_attack, mixed_seed_groups): """Test that AtomicAttack handles mixed seed groups (some with user_messages, some without).""" atomic_attack = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=mixed_seed_groups, atomic_attack_name="Mixed Attack", ) assert len(atomic_attack.seed_groups) == 2 # First has no user_messages (empty list) assert len(atomic_attack.seed_groups[0].user_messages) == 0 # Second has user_messages assert len(atomic_attack.seed_groups[1].user_messages) == 2 @pytest.mark.usefixtures("patch_central_database") class TestEnrichAtomicAttackIdentifiers: """Tests for _enrich_atomic_attack_identifiers in AtomicAttack.""" async def test_enrichment_populates_atomic_attack_identifier(self, mock_attack): """Test that run_async enriches results with atomic_attack_identifier.""" seed_groups = [ AttackSeedGroup( seeds=[ SeedObjective(value="obj1"), SeedPrompt(value="technique1", is_general_technique=True), ] ), ] attack_id = ComponentIdentifier(class_name="MockAttack", class_module="test.mock") attack_result = AttackResult( conversation_id="conv-1", objective="obj1", outcome=AttackOutcome.SUCCESS, executed_turns=1, atomic_attack_identifier=AtomicAttackIdentifier.build(attack_identifier=attack_id), ) atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=seed_groups, atomic_attack_name="test" ) with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results([attack_result]) result = await atomic.run_async() enriched = result.completed_results[0] assert enriched.atomic_attack_identifier is not None assert enriched.atomic_attack_identifier.class_name == "AtomicAttack" assert "attack_technique" in enriched.atomic_attack_identifier.children assert "seed_identifiers" in enriched.atomic_attack_identifier.children async def test_enrichment_populates_even_when_result_has_no_prior_identifier(self, mock_attack): """Test that enrichment works even when result has no prior atomic_attack_identifier, since AttackTechnique.get_identifier() is self-contained.""" seed_groups = [ AttackSeedGroup(seeds=[SeedObjective(value="obj1"), SeedPrompt(value="p1")]), ] attack_result = AttackResult( conversation_id="conv-1", objective="obj1", outcome=AttackOutcome.SUCCESS, executed_turns=1, atomic_attack_identifier=None, ) atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=seed_groups, atomic_attack_name="test" ) with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results([attack_result]) result = await atomic.run_async() # Should be enriched — technique provides its own identifier enriched = result.completed_results[0] assert enriched.atomic_attack_identifier is not None assert enriched.atomic_attack_identifier.class_name == "AtomicAttack" async def test_enrichment_skips_out_of_range_index(self, mock_attack): """Test that enrichment is skipped when input_indices has an out-of-range value.""" seed_groups = [ AttackSeedGroup(seeds=[SeedObjective(value="obj1"), SeedPrompt(value="p1")]), ] attack_id = ComponentIdentifier(class_name="MockAttack", class_module="test.mock") attack_result = AttackResult( conversation_id="conv-1", objective="obj1", outcome=AttackOutcome.SUCCESS, executed_turns=1, atomic_attack_identifier=AtomicAttackIdentifier.build(attack_identifier=attack_id), ) atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=seed_groups, atomic_attack_name="test" ) with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: # Index 99 is out of range for seed_groups (only 1 element) mock_exec.return_value = AttackExecutorResult( completed_results=[attack_result], incomplete_objectives=[], input_indices=[99], ) result = await atomic.run_async() # Should not be enriched (index out of range), so the identifier # should still lack seed info (seeds remains empty) enriched = result.completed_results[0] assert enriched.atomic_attack_identifier is not None seeds = enriched.atomic_attack_identifier.children.get("seeds", []) assert seeds == [], "Expected no seeds since index was out of range" async def test_enrichment_includes_all_seeds(self, mock_attack): """Test that all seeds (general and non-general) appear in the enriched identifier.""" seed_groups = [ AttackSeedGroup( seeds=[ SeedObjective(value="obj1"), SeedPrompt(value="technique", is_general_technique=True, value_sha256="tech_hash"), SeedPrompt(value="non_technique", is_general_technique=False, value_sha256="other_hash"), ] ), ] attack_id = ComponentIdentifier(class_name="MockAttack", class_module="test.mock") attack_result = AttackResult( conversation_id="conv-1", objective="obj1", outcome=AttackOutcome.SUCCESS, executed_turns=1, atomic_attack_identifier=AtomicAttackIdentifier.build(attack_identifier=attack_id), ) atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=seed_groups, atomic_attack_name="test" ) with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results([attack_result]) result = await atomic.run_async() enriched = result.completed_results[0].atomic_attack_identifier assert enriched is not None seed_ids = enriched.children["seed_identifiers"] # All three seeds (objective + technique + non_technique) should be present assert len(seed_ids) == 3 sha_values = [s.params.get("value_sha256") for s in seed_ids] assert "tech_hash" in sha_values assert "other_hash" in sha_values async def test_enrichment_maps_multiple_results_to_correct_seed_groups(self, mock_attack): """Test that multiple results are correctly mapped to their corresponding seed groups.""" seed_groups = [ AttackSeedGroup( seeds=[ SeedObjective(value="obj1"), SeedPrompt(value="tech_a", is_general_technique=True, value_sha256="hash_a"), ] ), AttackSeedGroup( seeds=[ SeedObjective(value="obj2"), SeedPrompt(value="tech_b", is_general_technique=True, value_sha256="hash_b"), ] ), ] attack_id = ComponentIdentifier(class_name="MockAttack", class_module="test.mock") results = [ AttackResult( conversation_id="c1", objective="obj1", outcome=AttackOutcome.SUCCESS, executed_turns=1, atomic_attack_identifier=AtomicAttackIdentifier.build(attack_identifier=attack_id), ), AttackResult( conversation_id="c2", objective="obj2", outcome=AttackOutcome.SUCCESS, executed_turns=1, atomic_attack_identifier=AtomicAttackIdentifier.build(attack_identifier=attack_id), ), ] atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=seed_groups, atomic_attack_name="test" ) with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results(results) result = await atomic.run_async() # First result should have hash_a seed enriched_0 = result.completed_results[0].atomic_attack_identifier seed_sha_values_0 = [s.params.get("value_sha256") for s in enriched_0.children["seed_identifiers"]] assert "hash_a" in seed_sha_values_0 # Second result should have hash_b seed enriched_1 = result.completed_results[1].atomic_attack_identifier seed_sha_values_1 = [s.params.get("value_sha256") for s in enriched_1.children["seed_identifiers"]] assert "hash_b" in seed_sha_values_1 async def test_enrichment_persists_to_db(self, mock_attack): """Test that enrichment persists the updated atomic_attack_identifier to the database.""" seed_groups = [ AttackSeedGroup( seeds=[ SeedObjective(value="obj1"), SeedPrompt(value="technique1", is_general_technique=True), ] ), ] attack_id = ComponentIdentifier(class_name="MockAttack", class_module="test.mock") attack_result = AttackResult( conversation_id="conv-1", objective="obj1", outcome=AttackOutcome.SUCCESS, executed_turns=1, attack_result_id="00000000-0000-0000-0000-000000000001", atomic_attack_identifier=AtomicAttackIdentifier.build(attack_identifier=attack_id), ) atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=seed_groups, atomic_attack_name="test" ) with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results([attack_result]) mock_memory = MagicMock() mock_memory.update_attack_result_by_id.return_value = True with patch("pyrit.scenario.core.atomic_attack.CentralMemory") as mock_cm: mock_cm.get_memory_instance.return_value = mock_memory await atomic.run_async() mock_memory.update_attack_result_by_id.assert_called_once() call_kwargs = mock_memory.update_attack_result_by_id.call_args.kwargs assert call_kwargs["attack_result_id"] == "00000000-0000-0000-0000-000000000001" assert "atomic_attack_identifier" in call_kwargs["update_fields"] # The persisted dict should have the AtomicAttack shape persisted = call_kwargs["update_fields"]["atomic_attack_identifier"] assert persisted["class_name"] == "AtomicAttack" async def test_enrichment_skips_db_update_when_no_attack_result_id(self, mock_attack): """Test that enrichment does not attempt a DB update when attack_result_id is empty.""" seed_groups = [ AttackSeedGroup( seeds=[ SeedObjective(value="obj1"), SeedPrompt(value="technique1", is_general_technique=True), ] ), ] attack_id = ComponentIdentifier(class_name="MockAttack", class_module="test.mock") attack_result = AttackResult( conversation_id="conv-1", objective="obj1", outcome=AttackOutcome.SUCCESS, executed_turns=1, attack_result_id="", atomic_attack_identifier=AtomicAttackIdentifier.build(attack_identifier=attack_id), ) atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=seed_groups, atomic_attack_name="test" ) with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results([attack_result]) mock_memory = MagicMock() with patch("pyrit.scenario.core.atomic_attack.CentralMemory") as mock_cm: mock_cm.get_memory_instance.return_value = mock_memory await atomic.run_async() mock_memory.update_attack_result_by_id.assert_not_called() @pytest.mark.usefixtures("patch_central_database") class TestAtomicAttackFilterSeedGroupsByCompletedHashes: """Tests for ``drop_seed_groups_with_hashes`` — the hash-based resume filter.""" def test_filters_out_completed_hashes(self, mock_attack, sample_seed_groups): from pyrit.common.utils import to_sha256 atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="test", ) completed = {to_sha256("objective1"), to_sha256("objective3")} atomic.drop_seed_groups_with_hashes(hashes=completed) assert atomic.seed_groups == [sample_seed_groups[1]] def test_empty_completed_hashes_is_noop(self, mock_attack, sample_seed_groups): atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="test", ) atomic.drop_seed_groups_with_hashes(hashes=set()) assert atomic.seed_groups == sample_seed_groups def test_all_hashes_completed_clears_seed_groups(self, mock_attack, sample_seed_groups): from pyrit.common.utils import to_sha256 atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="test", ) atomic.drop_seed_groups_with_hashes(hashes={to_sha256(f"objective{i}") for i in range(1, 4)}) assert atomic.seed_groups == [] def test_filter_is_stable_across_resampling(self, mock_attack, sample_seed_groups): """Identity is content-derived, so reordering ``_seed_groups`` between two calls (e.g. a fresh ``random.sample``) doesn't break the filter.""" from pyrit.common.utils import to_sha256 atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="test", ) # Simulate a re-sample by reversing the internal list. atomic._seed_groups = list(reversed(atomic._seed_groups)) atomic.drop_seed_groups_with_hashes(hashes={to_sha256("objective1")}) kept_objectives = [sg.objective.value for sg in atomic.seed_groups] assert "objective1" not in kept_objectives assert set(kept_objectives) == {"objective2", "objective3"} @pytest.mark.usefixtures("patch_central_database") class TestAtomicAttackRestrictSeedGroupsToHashes: """Tests for ``keep_seed_groups_with_hashes`` — the keep-set inverse used on resume to replay the originally-sampled subset.""" def test_keeps_only_listed_hashes(self, mock_attack, sample_seed_groups): from pyrit.common.utils import to_sha256 atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="test", ) keep = {to_sha256("objective1"), to_sha256("objective3")} retained = atomic.keep_seed_groups_with_hashes(hashes=keep) assert {sg.objective.value for sg in atomic.seed_groups} == {"objective1", "objective3"} assert retained == keep def test_retained_set_excludes_missing_hashes(self, mock_attack, sample_seed_groups): from pyrit.common.utils import to_sha256 atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="test", ) keep = {to_sha256("objective1"), to_sha256("not-in-dataset")} retained = atomic.keep_seed_groups_with_hashes(hashes=keep) assert {sg.objective.value for sg in atomic.seed_groups} == {"objective1"} assert retained == {to_sha256("objective1")} @pytest.mark.usefixtures("patch_central_database") class TestAtomicAttackDuplicateObjectiveValidation: """``AtomicAttack.__init__`` enforces objective-hash uniqueness within a single atomic attack so resume can use the hash as a stable identity.""" def test_constructing_with_duplicate_objective_raises(self, mock_attack): duplicate_groups = [ AttackSeedGroup(seeds=[SeedObjective(value="same-objective")]), AttackSeedGroup(seeds=[SeedObjective(value="same-objective")]), ] with pytest.raises(ValueError, match="duplicate objective hash"): AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=duplicate_groups, atomic_attack_name="dup", ) def test_constructing_with_unique_objectives_succeeds(self, mock_attack, sample_seed_groups): atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="ok", ) assert len(atomic.seed_groups) == 3 @pytest.mark.usefixtures("patch_central_database") class TestAtomicAttackAttributionStamping: """Tests for how ``run_async`` builds the ``AttackResultAttribution`` it passes to the executor.""" async def test_no_attribution_when_scenario_result_id_unset( self, mock_attack, sample_seed_groups, sample_attack_results ): """Outside a Scenario, ``_scenario_result_id`` is None and the executor must receive ``attributions=None``.""" atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="test", ) assert atomic._scenario_result_id is None with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results(sample_attack_results) await atomic.run_async() assert mock_exec.call_args.kwargs["attributions"] is None async def test_attribution_built_when_scenario_result_id_set( self, mock_attack, sample_seed_groups, sample_attack_results ): """When the Scenario stamps ``_scenario_result_id`` onto the atomic attack, ``run_async`` must build and pass per-seed-group attribution.""" from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="MyAtomicAttack", ) atomic._scenario_result_id = "00000000-0000-0000-0000-000000000abc" with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results(sample_attack_results) await atomic.run_async() attributions = mock_exec.call_args.kwargs["attributions"] assert len(attributions) == len(sample_seed_groups) assert all(isinstance(attribution, AttackResultAttribution) for attribution in attributions) assert all(attribution.parent_id == "00000000-0000-0000-0000-000000000abc" for attribution in attributions) assert all(attribution.parent_collection == "MyAtomicAttack" for attribution in attributions) assert [attribution.seed_group_id for attribution in attributions] == [ seed_group.logical_id for seed_group in sample_seed_groups ] async def test_attribution_includes_technique_eval_hash( self, mock_attack, sample_seed_groups, sample_attack_results ): """The stamped attribution must carry ``parent_eval_hash`` equal to ``technique_eval_hash`` so resume disambiguates between two atomic attacks that share a name but use different techniques.""" atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="MyAtomicAttack", ) atomic._scenario_result_id = "00000000-0000-0000-0000-000000000abc" with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.return_value = wrap_results(sample_attack_results) await atomic.run_async() attributions = mock_exec.call_args.kwargs["attributions"] assert all(attribution.parent_eval_hash is not None for attribution in attributions) assert all(attribution.parent_eval_hash == atomic.technique_eval_hash for attribution in attributions) @pytest.mark.usefixtures("patch_central_database") class TestAtomicAttackTechniqueEvalHash: """``technique_eval_hash`` must be stable across seed groups and differ between distinct technique configurations — it's the resume bucket key.""" def test_hash_is_independent_of_seed_groups(self, mock_attack, sample_seed_groups): a1 = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="same", ) a2 = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=[AttackSeedGroup(seeds=[SeedObjective(value="different-objective")])], atomic_attack_name="same", ) assert a1.technique_eval_hash == a2.technique_eval_hash def test_hash_differs_for_different_attacks(self, sample_seed_groups): attack_a = MagicMock(spec=AttackStrategy) attack_a.get_identifier.return_value = ComponentIdentifier(class_name="AttackA", class_module="pyrit.test") attack_b = MagicMock(spec=AttackStrategy) attack_b.get_identifier.return_value = ComponentIdentifier(class_name="AttackB", class_module="pyrit.test") a1 = AtomicAttack( attack_technique=AttackTechnique(attack=attack_a), seed_groups=sample_seed_groups, atomic_attack_name="same", ) a2 = AtomicAttack( attack_technique=AttackTechnique(attack=attack_b), seed_groups=sample_seed_groups, atomic_attack_name="same", ) assert a1.technique_eval_hash != a2.technique_eval_hash