mirror of
https://github.com/msoedov/agentic_security.git
synced 2026-06-29 16:49:58 +02:00
feat(restruct tests):
This commit is contained in:
Alexander Myasoedov
@@ -0,0 +1,285 @@
|
||||
import asyncio
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from agentic_security.primitives import Scan
|
||||
from agentic_security.probe_actor.fuzzer import (
|
||||
FuzzerState,
|
||||
generate_prompts,
|
||||
perform_many_shot_scan,
|
||||
perform_single_shot_scan,
|
||||
process_prompt,
|
||||
scan_router,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_prompts_with_list():
|
||||
prompts = ["prompt1", "prompt2", "prompt3"]
|
||||
results = [p async for p in generate_prompts(prompts)]
|
||||
assert results == prompts
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_prompts_with_async_generator():
|
||||
async def async_gen():
|
||||
for i in range(3):
|
||||
yield f"prompt{i}"
|
||||
|
||||
results = [p async for p in generate_prompts(async_gen())]
|
||||
assert results == ["prompt0", "prompt1", "prompt2"]
|
||||
|
||||
|
||||
async def assert_scan(generator, messages):
|
||||
results = [r async for r in generator]
|
||||
|
||||
for m in messages:
|
||||
found = False
|
||||
for r in results:
|
||||
if m in r:
|
||||
found = True
|
||||
break
|
||||
assert found, f"Message '{m}' not found in results. Results: {results}"
|
||||
return results
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("agentic_security.probe_data.data.prepare_prompts")
|
||||
async def test_perform_single_shot_scan_success(prepare_prompts_mock):
|
||||
# Mock prompt modules
|
||||
prepare_prompts_mock.return_value = [
|
||||
MagicMock(
|
||||
dataset_name="test_module",
|
||||
prompts=["test_prompt1", "test_prompt2"],
|
||||
lazy=False,
|
||||
)
|
||||
]
|
||||
|
||||
# Mock request_factory
|
||||
mock_response = AsyncMock()
|
||||
mock_response.fn.return_value = AsyncMock(
|
||||
status_code=200, text="response text", json=lambda: {}
|
||||
)
|
||||
request_factory = mock_response
|
||||
|
||||
async_gen = perform_single_shot_scan(
|
||||
request_factory=request_factory,
|
||||
max_budget=100,
|
||||
datasets=[{"dataset_name": "test", "selected": True}],
|
||||
optimize=False,
|
||||
)
|
||||
|
||||
await assert_scan(async_gen, ["Loading", "Scan completed."])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("agentic_security.probe_data.msj_data.prepare_prompts")
|
||||
@patch("agentic_security.probe_data.data.prepare_prompts")
|
||||
async def test_perform_many_shot_scan_probe_injection(
|
||||
prepare_prompts_mock, msj_prepare_prompts_mock
|
||||
):
|
||||
# Mock main and probe prompt modules
|
||||
prepare_prompts_mock.side_effect = [
|
||||
[MagicMock(dataset_name="main_module", prompts=["main_prompt1"], lazy=False)],
|
||||
[MagicMock(dataset_name="probe_module", prompts=["probe_prompt1"], lazy=False)],
|
||||
]
|
||||
|
||||
msj_prepare_prompts_mock.return_value = [
|
||||
MagicMock(
|
||||
dataset_name="msj_probe_module", prompts=["msj_probe_prompt"], lazy=False
|
||||
)
|
||||
]
|
||||
|
||||
# Mock request_factory
|
||||
mock_response = AsyncMock()
|
||||
mock_response.fn.side_effect = [
|
||||
AsyncMock(status_code=200, text="main response", json=lambda: {}),
|
||||
AsyncMock(status_code=200, text="probe response", json=lambda: {}),
|
||||
]
|
||||
request_factory = mock_response
|
||||
|
||||
async_gen = perform_many_shot_scan(
|
||||
request_factory=request_factory,
|
||||
max_budget=100,
|
||||
datasets=[{"dataset_name": "main", "selected": True}],
|
||||
probe_datasets=[{"dataset_name": "probe", "selected": True}],
|
||||
probe_frequency=1.0, # Always inject probes
|
||||
optimize=False,
|
||||
)
|
||||
|
||||
await assert_scan(async_gen, ["Loading", "Scan completed."])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("agentic_security.probe_data.data.prepare_prompts")
|
||||
async def test_scan_router_single_shot(prepare_prompts_mock):
|
||||
prepare_prompts_mock.return_value = []
|
||||
|
||||
request_factory = AsyncMock()
|
||||
scan_params = Scan(
|
||||
maxBudget=100,
|
||||
llmSpec="test",
|
||||
datasets=[],
|
||||
probe_datasets=[],
|
||||
enableMultiStepAttack=False,
|
||||
optimize=False,
|
||||
)
|
||||
|
||||
gen = scan_router(
|
||||
request_factory=request_factory,
|
||||
scan_parameters=scan_params,
|
||||
)
|
||||
await assert_scan(gen, ["Loading", "Scan completed."])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("agentic_security.probe_data.data.prepare_prompts")
|
||||
async def test_scan_router_many_shot(prepare_prompts_mock):
|
||||
prepare_prompts_mock.return_value = []
|
||||
|
||||
request_factory = AsyncMock()
|
||||
scan_params = Scan(
|
||||
maxBudget=100,
|
||||
datasets=[],
|
||||
llmSpec="test",
|
||||
probeDatasets=[],
|
||||
enableMultiStepAttack=True,
|
||||
optimize=False,
|
||||
)
|
||||
|
||||
gen = scan_router(
|
||||
request_factory=request_factory,
|
||||
scan_parameters=scan_params,
|
||||
)
|
||||
assert gen is not None
|
||||
|
||||
await assert_scan(gen, ["Loading", "Scan completed."])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_perform_single_shot_scan_stop_event():
|
||||
stop_event = asyncio.Event()
|
||||
stop_event.set() # Pre-set to simulate user stopping the scan
|
||||
|
||||
async def request_mock(*args, **kwargs):
|
||||
return AsyncMock(status_code=200, text="response text", json=lambda: {})
|
||||
|
||||
async_gen = perform_single_shot_scan(
|
||||
request_factory=MagicMock(fn=request_mock),
|
||||
max_budget=100,
|
||||
datasets=[],
|
||||
stop_event=stop_event,
|
||||
)
|
||||
|
||||
await assert_scan(async_gen, ["Loading", "Scan completed."])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_perform_many_shot_scan_stop_event():
|
||||
stop_event = asyncio.Event()
|
||||
stop_event.set() # Pre-set to simulate user stopping the scan
|
||||
|
||||
async def request_mock(*args, **kwargs):
|
||||
return AsyncMock(status_code=200, text="response text", json=lambda: {})
|
||||
|
||||
async_gen = perform_many_shot_scan(
|
||||
request_factory=MagicMock(fn=request_mock),
|
||||
max_budget=100,
|
||||
datasets=[],
|
||||
probe_datasets=[],
|
||||
stop_event=stop_event,
|
||||
)
|
||||
|
||||
await assert_scan(async_gen, ["Loading", "Scan completed."])
|
||||
|
||||
|
||||
def mock_refusal_heuristic(response_json):
|
||||
return response_json.get("is_refusal", False)
|
||||
|
||||
|
||||
class TestProcessPrompt(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_successful_response_no_refusal(self):
|
||||
mock_request_factory = Mock()
|
||||
mock_request_factory.fn = AsyncMock(
|
||||
return_value=Mock(
|
||||
status_code=200,
|
||||
text="Valid response text",
|
||||
json=Mock(return_value={"is_refusal": False}),
|
||||
request="mock_request",
|
||||
)
|
||||
)
|
||||
|
||||
tokens, refusal = await process_prompt(
|
||||
request_factory=mock_request_factory,
|
||||
prompt="test prompt",
|
||||
tokens=0,
|
||||
module_name="module_a",
|
||||
fuzzer_state=FuzzerState(),
|
||||
)
|
||||
|
||||
self.assertEqual(tokens, 3) # Tokens from "Valid response text"
|
||||
self.assertTrue(refusal)
|
||||
|
||||
async def test_successful_response_with_refusal(self):
|
||||
mock_request_factory = Mock()
|
||||
mock_request_factory.fn = AsyncMock(
|
||||
return_value=Mock(
|
||||
status_code=200,
|
||||
text="Response indicating refusal",
|
||||
json=Mock(return_value={"is_refusal": True}),
|
||||
request="mock_request",
|
||||
)
|
||||
)
|
||||
|
||||
fuzzer_state = FuzzerState()
|
||||
tokens, refusal = await process_prompt(
|
||||
request_factory=mock_request_factory,
|
||||
prompt="test prompt",
|
||||
tokens=0,
|
||||
module_name="module_a",
|
||||
fuzzer_state=fuzzer_state,
|
||||
)
|
||||
|
||||
self.assertEqual(tokens, 3) # Tokens from "Response indicating refusal"
|
||||
# self.assertFalse(fuzzer_state.refusals)
|
||||
|
||||
async def test_http_error_response(self):
|
||||
mock_request_factory = Mock()
|
||||
mock_request_factory.fn = AsyncMock(
|
||||
return_value=Mock(
|
||||
status_code=500,
|
||||
text="Internal Server Error",
|
||||
request="mock_request",
|
||||
response=Mock(),
|
||||
)
|
||||
)
|
||||
|
||||
fuzzer_state = FuzzerState()
|
||||
await process_prompt(
|
||||
request_factory=mock_request_factory,
|
||||
prompt="test prompt",
|
||||
tokens=0,
|
||||
module_name="module_a",
|
||||
fuzzer_state=fuzzer_state,
|
||||
)
|
||||
|
||||
async def test_request_error(self):
|
||||
mock_request_factory = Mock()
|
||||
mock_request_factory.fn = AsyncMock(
|
||||
side_effect=httpx.RequestError("Connection error")
|
||||
)
|
||||
|
||||
fuzzer_state = FuzzerState()
|
||||
tokens, refusal = await process_prompt(
|
||||
request_factory=mock_request_factory,
|
||||
prompt="test prompt",
|
||||
tokens=0,
|
||||
module_name="module_a",
|
||||
fuzzer_state=fuzzer_state,
|
||||
)
|
||||
|
||||
self.assertEqual(tokens, 0)
|
||||
self.assertTrue(refusal)
|
||||
@@ -0,0 +1,13 @@
|
||||
from agentic_security.probe_actor.refusal import DefaultRefusalClassifier
|
||||
|
||||
|
||||
class TestCheckRefusal:
|
||||
# The function correctly identifies a refusal phrase in the response.
|
||||
def test_identify_refusal_phrase(self):
|
||||
response = "I'm sorry, but I cannot provide that information."
|
||||
assert DefaultRefusalClassifier().is_refusal(response)
|
||||
|
||||
# The response is an empty string.
|
||||
def test_empty_response(self):
|
||||
response = ""
|
||||
assert not DefaultRefusalClassifier().is_refusal(response)
|
||||
Reference in New Issue
Block a user