diff --git a/doc/code/framework.md b/doc/code/framework.md index e046041bdc..766827b74d 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -185,6 +185,7 @@ If you are contributing to PyRIT, that work will most likely land in one of the - `core` stays deliberately small so a default run doesn't print 200 techniques or take forever; the wider catalog lives in `extra` and is selected on demand. Users pick subsets by passing initializer tags (e.g. `core`, `extra`, `all`) or writing their own initializer, so different runs — including from the CLI — can register different technique sets without changing the catalog. - A technique tied to one scenario is fine; if it's pinned and non-reusable it can stay local to that scenario, but if another scenario could reuse it, promote it to a catalog module and tag it. - Tags describe a technique (behavioral tags like `single_turn`/`multi_turn`, owner tags like `airt`); they don't decide what a scenario runs. There is deliberately **no global `default` tag** — a default is scenario-relative, declared per scenario via `build_technique_class_from_factories` (the `factories` list is the pool, catalog tags become named aggregate presets, and `default_tags` / `default_names` set what runs when nothing is chosen). +- Generated technique enums follow their factory pool's lifecycle. Registry-derived enums and their lazy exports are discarded on setup registry reset so new scenarios see the replacement catalog. Enums built exclusively from fixed scenario-owned factories, such as `DoctorTechnique` and `JailbreakTechnique`, retain their identity across registry resets, so previously imported members remain valid for new scenario instances. - Factories opt into additive request-converter composition with `supports_additional_request_converters=True`. This is a semantic capability, not just constructor-signature detection; the factory validates that opted-in attacks accept `attack_converter_config`. - **Does not own**: the conversation algorithm itself. Branching, turn management, and scoring decisions live in the executor it wraps — a technique only selects and configures existing components, and shouldn't implement new sending, scoring, or branching logic. diff --git a/doc/contributing/4_running_tests.md b/doc/contributing/4_running_tests.md index f052c57e69..20f862aa9c 100644 --- a/doc/contributing/4_running_tests.md +++ b/doc/contributing/4_running_tests.md @@ -33,6 +33,11 @@ prints the 25 slowest test phases. Override the report path with CI uploads JUnit reports for every OS, Python version, and extras combination, including failed test runs, with those dimensions in each artifact name. +LiteLLM unit tests mock the SDK and do not need tokenizer data or a local HTTP +server, including when the `litellm` or `all` extra is installed. Real adapter +and outgoing-header checks run separately as +[loopback integration tests](./6_integration_tests.md#litellm-loopback-integration-tests). + ## Running a subset while iterating For a narrower run, invoke `pytest` directly. You can invoke pytest if it's in your path or via python; either `pytest` or `python -m pytest`. For the following examples, we will use `pytest`. diff --git a/doc/contributing/6_integration_tests.md b/doc/contributing/6_integration_tests.md index bbbb74d4f7..663da84ce6 100644 --- a/doc/contributing/6_integration_tests.md +++ b/doc/contributing/6_integration_tests.md @@ -8,3 +8,42 @@ These are tested regularly by the PyRIT team but not necessarily run every pull - Unit tests should test all scenarios and more is often better. Integration tests should be scoped and be careful to not run too long. - Integration tests can sometimes test end-to-end. But almost always, these should target one scenario. + +## LiteLLM loopback integration tests + +`tests/integration/targets/test_litellm_trace_integration.py` checks the real +OpenAI and Anthropic LiteLLM adapters against a local HTTP server. It verifies +the actual outgoing trace and routing headers and response parsing. These tests +use a fake API key and isolated SQLite memory, not live provider credentials. +Like other integration tests, they are gated by `RUN_ALL_TESTS=true` and need the +optional `litellm` dependency. + +Anthropic token counting requires the official +[`cl100k_base.tiktoken`](https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken) +data, which is not included in the tiktoken wheel. Obtain a local copy separately +before offline execution and point the integration fixture to it: + +```powershell +$env:PYRIT_TEST_TIKTOKEN_ASSET = 'C:\test-inputs\cl100k_base.tiktoken' +$env:RUN_ALL_TESTS = 'true' +uv run -m pytest tests\integration\targets\test_litellm_trace_integration.py +``` + +On a POSIX shell: + +```bash +PYRIT_TEST_TIKTOKEN_ASSET=/tmp/test-inputs/cl100k_base.tiktoken RUN_ALL_TESTS=true \ + uv run -m pytest tests/integration/targets/test_litellm_trace_integration.py +``` + +The fixture checks SHA-256 +`223921b76ee99bde995b7ff738513eef100fb51d18c93597a113bcffe865b2a7`, +creates a private temporary tokenizer cache, and loads the real encoding without +relying on a warm process cache. It selects LiteLLM's bundled model-cost and +Anthropic-header metadata through supported local-only flags. It never downloads +data. Missing or corrupt tokenizer input is an explicit preparation error once +the tests are enabled, not a skipped test. + +Keep the data outside Git. Do not disable offline guards or allow downloads +inside pytest. Ordinary unit runs need neither this asset nor the integration +environment variables. diff --git a/pyrit/scenario/scenarios/_dynamic_techniques.py b/pyrit/scenario/scenarios/_dynamic_techniques.py index 5d699a4151..7bef15b7ce 100644 --- a/pyrit/scenario/scenarios/_dynamic_techniques.py +++ b/pyrit/scenario/scenarios/_dynamic_techniques.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -"""Lazy exports for scenario technique classes built from registered catalogs.""" +"""Lazy exports for scenario technique classes built from registered or local factories.""" from importlib import import_module from typing import TYPE_CHECKING @@ -17,14 +17,12 @@ MultilingualTechnique: type[ScenarioTechnique] RapidResponseTechnique: type[ScenarioTechnique] -_TECHNIQUE_BUILDERS = { +_REGISTRY_TECHNIQUE_BUILDERS = { "AdversarialBenchmarkTechnique": ( "pyrit.scenario.scenarios.benchmark.adversarial", "_build_benchmark_technique", ), "CyberTechnique": ("pyrit.scenario.scenarios.airt.cyber", "_build_cyber_technique"), - "DoctorTechnique": ("pyrit.scenario.scenarios.garak.doctor", "_build_doctor_technique"), - "JailbreakTechnique": ("pyrit.scenario.scenarios.airt.jailbreak", "_build_jailbreak_technique"), "LeakageTechnique": ("pyrit.scenario.scenarios.airt.leakage", "_build_leakage_technique"), "MultilingualTechnique": ( "pyrit.scenario.scenarios.airt.multilingual", @@ -36,6 +34,13 @@ ), } +_TECHNIQUE_BUILDERS = { + **_REGISTRY_TECHNIQUE_BUILDERS, + # Fixed scenario-owned pools must keep enum identity across unrelated registry resets. + "DoctorTechnique": ("pyrit.scenario.scenarios.garak.doctor", "_build_doctor_technique"), + "JailbreakTechnique": ("pyrit.scenario.scenarios.airt.jailbreak", "_build_jailbreak_technique"), +} + def __getattr__(name: str) -> object: try: @@ -54,7 +59,7 @@ def reset_dynamic_technique_caches() -> None: "pyrit.scenario.scenarios.benchmark", "pyrit.scenario.scenarios.garak", } - for technique_name, (module_name, builder_name) in _TECHNIQUE_BUILDERS.items(): + for technique_name, (module_name, builder_name) in _REGISTRY_TECHNIQUE_BUILDERS.items(): builder = getattr(import_module(module_name), builder_name) builder.cache_clear() globals().pop(technique_name, None) diff --git a/tests/integration/targets/test_litellm_trace_integration.py b/tests/integration/targets/test_litellm_trace_integration.py new file mode 100644 index 0000000000..0142085e62 --- /dev/null +++ b/tests/integration/targets/test_litellm_trace_integration.py @@ -0,0 +1,186 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Verify trace headers through real LiteLLM adapters against a loopback server.""" + +import hashlib +import json +import os +import threading +from collections.abc import Callable, Iterator +from email.message import Message as HTTPHeaders +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import TYPE_CHECKING +from unittest.mock import patch + +import pytest + +from pyrit.models import Message, MessagePiece, RequestTraceContext +from pyrit.prompt_target import LiteLLMChatTarget, TargetTraceConfig + +if TYPE_CHECKING: + import tiktoken + +pytestmark = [pytest.mark.run_only_if_all_tests, pytest.mark.usefixtures("patch_central_database")] + +_TOKENIZER_URL = "https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken" +_TOKENIZER_SHA256 = "223921b76ee99bde995b7ff738513eef100fb51d18c93597a113bcffe865b2a7" +_OPENAI_REPLY = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 0, + "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hello"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} +_ANTHROPIC_REPLY = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, +} +_WIRE_PROVIDERS = pytest.mark.parametrize( + ("model_name", "base_path", "request_path"), + [ + ("openai/gpt-4o", "/v1", "/v1/chat/completions"), + ("anthropic/claude-sonnet-4-6", "", "/v1/messages"), + ], +) + + +@pytest.fixture +def official_cl100k_encoding(tmp_path: Path) -> Iterator["tiktoken.Encoding"]: + cache_dir = tmp_path / "official-tokenizer" + with patch.dict( + os.environ, + { + "TIKTOKEN_CACHE_DIR": str(cache_dir), + "LITELLM_LOCAL_MODEL_COST_MAP": "true", + "LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS": "true", + }, + ): + pytest.importorskip("litellm") + import tiktoken + import tiktoken.registry + + source = os.environ.get("PYRIT_TEST_TIKTOKEN_ASSET") + if not source: + raise RuntimeError( + "Set PYRIT_TEST_TIKTOKEN_ASSET to the official cl100k_base.tiktoken file. " + "See doc/contributing/6_integration_tests.md for offline preparation." + ) + content = Path(source).read_bytes() + if hashlib.sha256(content).hexdigest() != _TOKENIZER_SHA256: + raise ValueError(f"Official cl100k_base asset SHA-256 mismatch: {source}") + cache_dir.mkdir() + (cache_dir / hashlib.sha1(_TOKENIZER_URL.encode("utf-8")).hexdigest()).write_bytes(content) + with patch.dict(tiktoken.registry.ENCODINGS, {}, clear=True): + yield tiktoken.get_encoding("cl100k_base") + + +@pytest.fixture +def chat_endpoint( + official_cl100k_encoding: "tiktoken.Encoding", +) -> Iterator[tuple[str, list[tuple[str, HTTPHeaders]]]]: + received: list[tuple[str, HTTPHeaders]] = [] + + class _Handler(BaseHTTPRequestHandler): + def do_POST(self) -> None: + self.rfile.read(int(self.headers.get("Content-Length", 0))) + received.append((self.path, self.headers)) + reply = _ANTHROPIC_REPLY if self.path.endswith("/v1/messages") else _OPENAI_REPLY + payload = json.dumps(reply).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, format: str, *args: object) -> None: # noqa: A002 + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), _Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}", received + finally: + server.shutdown() + server.server_close() + thread.join() + + +def _user_message() -> Message: + return MessagePiece(role="user", conversation_id="convo", original_value="test prompt").to_message() + + +@_WIRE_PROVIDERS +async def test_recorded_trace_context_reaches_wire_through_litellm_async( + *, + chat_endpoint: tuple[str, list[tuple[str, HTTPHeaders]]], + model_name: str, + base_path: str, + request_path: str, +) -> None: + url, received = chat_endpoint + provider = model_name.split("/")[0] + target = LiteLLMChatTarget( + model_name=model_name, + endpoint=url + base_path, + api_key="test-key", + extra_body_parameters={ + "provider_specific_header": {"custom_llm_provider": provider, "extra_headers": {"X-Route": "a"}} + }, + trace_config=TargetTraceConfig(enabled=True), + ) + request = _user_message() + responses = await target.send_prompt_async(message=request) + + link = RequestTraceContext.from_metadata(request.get_piece().prompt_metadata) + assert link is not None + assert len(received) == 1 + path, headers = received[0] + assert path == request_path + assert headers.get_all("traceparent") == [link.traceparent] + assert headers.get_all("tracestate") is None + assert headers["X-Route"] == "a" + assert responses[0].get_value() == "hello" + + +def _single_scope(provider: str) -> object: + return {"custom_llm_provider": provider, "extra_headers": {"TraceParent": f"00-{'3' * 32}-{'4' * 16}-01"}} + + +def _list_scope(provider: str) -> object: + return [ + {"custom_llm_provider": "bedrock", "extra_headers": {"X-Other": "1"}}, + {"custom_llm_provider": f"bedrock, {provider}", "extra_headers": {"tracestate": "vendor=manual"}}, + ] + + +@_WIRE_PROVIDERS +@pytest.mark.parametrize("scoped_headers", [_single_scope, _list_scope]) +async def test_provider_specific_trace_header_is_rejected_through_litellm_async( + *, + chat_endpoint: tuple[str, list[tuple[str, HTTPHeaders]]], + model_name: str, + base_path: str, + request_path: str, + scoped_headers: Callable[[str], object], +) -> None: + url, received = chat_endpoint + target = LiteLLMChatTarget( + model_name=model_name, + endpoint=url + base_path, + api_key="test-key", + extra_body_parameters={"provider_specific_header": scoped_headers(model_name.split("/")[0])}, + trace_config=TargetTraceConfig(enabled=True), + ) + with pytest.raises(ValueError, match="Manual trace headers"): + await target.send_prompt_async(message=_user_message()) + assert received == [] diff --git a/tests/unit/auth/test_azure_auth.py b/tests/unit/auth/test_azure_auth.py index 1bd5ec5abc..0fe720af43 100644 --- a/tests/unit/auth/test_azure_auth.py +++ b/tests/unit/auth/test_azure_auth.py @@ -6,8 +6,8 @@ from unittest.mock import MagicMock, patch import pytest +from azure.core.credentials import AccessToken -from pyrit.auth.auth_config import REFRESH_TOKEN_BEFORE_MSEC from pyrit.auth.azure_auth import ( AzureAuth, ensure_async_token_provider, @@ -20,6 +20,7 @@ curr_epoch_time = int(time.time()) mock_token = "fake token" +_TOKEN_TIME = 1_700_000_000 @pytest.mark.parametrize("returns_awaitable", [False, True]) @@ -44,32 +45,38 @@ def is_speechsdk_installed(): return False -def test_get_token_on_init(): - with patch("azure.identity.AzureCliCredential.get_token") as mock_get_token: - mock_get_token.return_value = MagicMock(token=mock_token) +def test_get_token_on_init() -> None: + with patch("azure.identity.DefaultAzureCredential.get_token", return_value=AccessToken(mock_token, _TOKEN_TIME)): test_instance = AzureAuth(token_scope="https://mocked_endpoint.azure.com") assert test_instance.token == mock_token -def test_refresh_no_expiration(): - # Token not expired so not reset - with patch("azure.identity.AzureCliCredential.get_token") as mock_get_token: - mock_get_token.return_value = MagicMock( - token=mock_token, expires_on=curr_epoch_time + REFRESH_TOKEN_BEFORE_MSEC - ) +@pytest.mark.parametrize("elapsed_seconds", [0, 301]) +def test_refresh_no_expiration(elapsed_seconds: int) -> None: + with ( + patch( + "azure.identity.DefaultAzureCredential.get_token", + return_value=AccessToken(mock_token, _TOKEN_TIME + 3600), + ) as mock_get_token, + patch("pyrit.auth.azure_auth.time.time", return_value=_TOKEN_TIME + elapsed_seconds), + ): test_instance = AzureAuth(token_scope="https://mocked_endpoint.azure.com") token = test_instance.refresh_token() assert token == mock_token - mock_get_token.assert_called() + mock_get_token.assert_called_once() -def test_refresh_expiration(): - # Token expired and reset - with patch("azure.identity.AzureCliCredential.get_token") as mock_get_token: - mock_get_token.return_value = MagicMock(token=mock_token, expires_on=curr_epoch_time) +def test_refresh_expiration() -> None: + with ( + patch( + "azure.identity.DefaultAzureCredential.get_token", + side_effect=[AccessToken("expired", _TOKEN_TIME), AccessToken(mock_token, _TOKEN_TIME + 3600)], + ) as mock_get_token, + patch("pyrit.auth.azure_auth.time.time", return_value=_TOKEN_TIME), + ): test_instance = AzureAuth(token_scope="https://mocked_endpoint.azure.com") token = test_instance.refresh_token() - assert token + assert token == mock_token assert mock_get_token.call_count == 2 diff --git a/tests/unit/datasets/test_comic_jailbreak_dataset.py b/tests/unit/datasets/test_comic_jailbreak_dataset.py index 7d584c10e6..ec2a01da0a 100644 --- a/tests/unit/datasets/test_comic_jailbreak_dataset.py +++ b/tests/unit/datasets/test_comic_jailbreak_dataset.py @@ -152,13 +152,21 @@ async def test_fetch_dataset_authors(self): assert len(seed.authors) == 3 assert seed.groups == ["Singapore University of Technology and Design"] - async def test_fetch_dataset_missing_goal_raises(self): - mock_data = [{"Target": "Sure", "Behavior": "Test", "Category": "Test"}] + @pytest.mark.parametrize("missing_key", ["Goal", "Category"]) + async def test_fetch_dataset_missing_goal_raises(self, missing_key: str) -> None: + example = _make_example() + del example[missing_key] + mock_data = [example] loader = _ComicJailbreakDataset(templates=["article"]) - with patch.object(loader, "_fetch_from_url", return_value=mock_data): - with pytest.raises(ValueError, match="Missing keys"): + with ( + patch.object(loader, "_fetch_from_url", return_value=mock_data), + patch.object(loader, "_fetch_template_async", new_callable=AsyncMock, return_value="/fake/template.png"), + patch.object(loader, "_render_comic_async", new_callable=AsyncMock) as render, + ): + with pytest.raises(ValueError, match=f"Missing keys in example: {missing_key}"): await loader.fetch_dataset_async() + render.assert_not_awaited() async def test_fetch_dataset_empty_goal_skipped(self): mock_data = [_make_example(Goal=" ")] diff --git a/tests/unit/prompt_target/target/test_litellm_chat_target.py b/tests/unit/prompt_target/target/test_litellm_chat_target.py index 009be98b02..6c5bbf338b 100644 --- a/tests/unit/prompt_target/target/test_litellm_chat_target.py +++ b/tests/unit/prompt_target/target/test_litellm_chat_target.py @@ -5,11 +5,8 @@ import base64 import json import sys -import threading import types -from collections.abc import Callable, Iterator -from email.message import Message as HTTPHeaders -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from collections.abc import Callable from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -948,78 +945,23 @@ async def test_provider_specific_trace_headers_require_disabled_tracing_async( # --------------------------------------------------------------------------- -# Request tracing through the real LiteLLM adapter +# Provider-specific request tracing at the LiteLLM boundary # --------------------------------------------------------------------------- -_OPENAI_REPLY = { - "id": "chatcmpl-1", - "object": "chat.completion", - "created": 0, - "model": "gpt-4o", - "choices": [{"index": 0, "message": {"role": "assistant", "content": "hello"}, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, -} -_ANTHROPIC_REPLY = { - "id": "msg_1", - "type": "message", - "role": "assistant", - "model": "claude-sonnet-4-6", - "content": [{"type": "text", "text": "hello"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 1, "output_tokens": 1}, -} -# LiteLLM appends ``/chat/completions`` to an OpenAI base and ``/v1/messages`` to an Anthropic base. -_WIRE_PROVIDERS = pytest.mark.parametrize( - ("model_name", "base_path"), [("openai/gpt-4o", "/v1"), ("anthropic/claude-sonnet-4-6", "")] -) - - -@pytest.fixture -def chat_endpoint() -> Iterator[tuple[str, list[HTTPHeaders]]]: - """Serve OpenAI and Anthropic chat replies on loopback and record the headers of each request.""" - pytest.importorskip("litellm") - received: list[HTTPHeaders] = [] - - class _Handler(BaseHTTPRequestHandler): - def do_POST(self) -> None: - self.rfile.read(int(self.headers.get("Content-Length", 0))) - received.append(self.headers) - reply = _ANTHROPIC_REPLY if self.path.endswith("/v1/messages") else _OPENAI_REPLY - payload = json.dumps(reply).encode() - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(payload))) - self.end_headers() - self.wfile.write(payload) - - def log_message(self, format: str, *args: object) -> None: # noqa: A002 - pass - - server = ThreadingHTTPServer(("127.0.0.1", 0), _Handler) - thread = threading.Thread(target=server.serve_forever, daemon=True) - thread.start() - try: - yield f"http://127.0.0.1:{server.server_address[1]}", received - finally: - server.shutdown() - server.server_close() - thread.join() - -@_WIRE_PROVIDERS -async def test_recorded_trace_context_reaches_wire_through_litellm_async( - patch_central_database, chat_endpoint: tuple[str, list[HTTPHeaders]], model_name: str, base_path: str +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize("model_name", ["openai/gpt-4o", "anthropic/claude-sonnet-4-6"]) +async def test_recorded_trace_context_reaches_litellm_arguments_async( + *, litellm_stub: types.ModuleType, model_name: str ) -> None: - url, received = chat_endpoint + litellm_stub.acompletion.return_value = _mock_response(model=model_name) provider = model_name.split("/")[0] + scoped_headers = {"custom_llm_provider": provider, "extra_headers": {"X-Route": "a"}} target = LiteLLMChatTarget( model_name=model_name, - endpoint=url + base_path, + endpoint="https://example.invalid", api_key="test-key", - extra_body_parameters={ - "provider_specific_header": {"custom_llm_provider": provider, "extra_headers": {"X-Route": "a"}} - }, + extra_body_parameters={"provider_specific_header": scoped_headers}, trace_config=TargetTraceConfig(enabled=True), ) request = _user_message() @@ -1027,10 +969,14 @@ async def test_recorded_trace_context_reaches_wire_through_litellm_async( link = RequestTraceContext.from_metadata(request.get_piece().prompt_metadata) assert link is not None - assert len(received) == 1 - assert received[0].get_all("traceparent") == [link.traceparent] - assert received[0].get_all("tracestate") is None - assert received[0]["X-Route"] == "a" + litellm_stub.acompletion.assert_awaited_once() + kwargs = litellm_stub.acompletion.call_args.kwargs + assert kwargs["model"] == model_name + assert kwargs["api_base"] == "https://example.invalid" + assert kwargs["api_key"] == "test-key" + assert kwargs["messages"] == [{"role": "user", "content": "test prompt"}] + assert kwargs["extra_headers"] == {"traceparent": link.traceparent} + assert kwargs["provider_specific_header"] == scoped_headers assert responses[0].get_value() == "hello" @@ -1045,26 +991,24 @@ def _list_scope(provider: str) -> object: ] -@_WIRE_PROVIDERS +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize("model_name", ["openai/gpt-4o", "anthropic/claude-sonnet-4-6"]) @pytest.mark.parametrize("scoped_headers", [_single_scope, _list_scope]) -async def test_provider_specific_trace_header_is_rejected_through_litellm_async( - patch_central_database, - chat_endpoint: tuple[str, list[HTTPHeaders]], +async def test_provider_specific_trace_header_is_rejected_before_litellm_async( + *, + litellm_stub: types.ModuleType, model_name: str, - base_path: str, scoped_headers: Callable[[str], object], ) -> None: - url, received = chat_endpoint target = LiteLLMChatTarget( model_name=model_name, - endpoint=url + base_path, api_key="test-key", extra_body_parameters={"provider_specific_header": scoped_headers(model_name.split("/")[0])}, trace_config=TargetTraceConfig(enabled=True), ) with pytest.raises(ValueError, match="Manual trace headers"): await target.send_prompt_async(message=_user_message()) - assert received == [] + litellm_stub.acompletion.assert_not_awaited() # --------------------------------------------------------------------------- diff --git a/tests/unit/prompt_target/test_discover_target_capabilities.py b/tests/unit/prompt_target/test_discover_target_capabilities.py index 510fb532f2..345de4ee94 100644 --- a/tests/unit/prompt_target/test_discover_target_capabilities.py +++ b/tests/unit/prompt_target/test_discover_target_capabilities.py @@ -866,6 +866,7 @@ def test_probe_removes_nested_and_legacy_tool_settings(target_type: type[PromptT model_name="unknown", endpoint="https://example.invalid", api_key="not-a-key", + custom_configuration=TargetConfiguration(capabilities=TargetCapabilities()), extra_body_parameters=body, ) with _permissive_configuration(target=target): diff --git a/tests/unit/scenario/garak/test_doctor.py b/tests/unit/scenario/garak/test_doctor.py index 7d459773bf..1b710b915a 100644 --- a/tests/unit/scenario/garak/test_doctor.py +++ b/tests/unit/scenario/garak/test_doctor.py @@ -3,6 +3,7 @@ """Tests for the Doctor scenario.""" +from collections.abc import Callable from unittest.mock import MagicMock, patch import pytest @@ -11,11 +12,19 @@ from pyrit.executor.attack import PromptSendingAttack from pyrit.models import ComponentIdentifier, SeedGroup, SeedObjective from pyrit.prompt_target import PromptTarget +from pyrit.registry import Registry from pyrit.scenario import DatasetAttackConfiguration from pyrit.scenario.core.scenario import BaselineAttackPolicy +from pyrit.scenario.core.scenario_technique import ScenarioTechnique from pyrit.scenario.garak import Doctor, DoctorTechnique # type: ignore[ty:unresolved-import] +from pyrit.scenario.scenarios._dynamic_techniques import reset_dynamic_technique_caches from pyrit.scenario.scenarios.garak.doctor import DOCTOR_FACTORIES from pyrit.score import TrueFalseScorer +from pyrit.setup.initialization import reset_setup_registries + + +class _OtherTechnique(ScenarioTechnique): + ALL = ("all", {"all"}) def _factories_by_name(): @@ -155,6 +164,37 @@ async def test_default_expands_to_concrete_techniques( technique_values = {s.value for s in scenario._scenario_techniques} assert technique_values == {"policy_puppetry", "policy_puppetry_leet"} + @pytest.mark.parametrize("reset_caches", [reset_dynamic_technique_caches, reset_setup_registries]) + async def test_imported_technique_survives_registry_reset_async( + self, + *, + mock_objective_target: PromptTarget, + mock_objective_scorer: TrueFalseScorer, + doctor_dataset_config: DatasetAttackConfiguration, + reset_caches: Callable[[], None], + ) -> None: + with patch.dict(Registry._singletons): + reset_caches() + scenario = Doctor(objective_scorer=mock_objective_scorer) + assert scenario._technique_class is DoctorTechnique + assert scenario._default_technique is DoctorTechnique.DEFAULT + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "scenario_techniques": [DoctorTechnique.ALL], + "dataset_config": doctor_dataset_config, + } + ) + await scenario.initialize_async() + + assert scenario._scenario_techniques == DoctorTechnique.get_all_techniques() + assert len(scenario._atomic_attacks) == 2 + + @pytest.mark.parametrize("unsupported", ["all", _OtherTechnique.ALL]) + def test_techniques_from_outside_doctor_catalog_are_rejected(self, unsupported: object) -> None: + with pytest.raises(ValueError, match="DoctorTechnique received unsupported techniques"): + DoctorTechnique.resolve([unsupported], default=DoctorTechnique.DEFAULT) + async def test_all_expands_to_concrete_techniques( self, mock_objective_target, mock_objective_scorer, doctor_dataset_config ): diff --git a/tests/unit/scenario/test_package_lazy_attrs.py b/tests/unit/scenario/test_package_lazy_attrs.py index a648c7044d..244077e4e5 100644 --- a/tests/unit/scenario/test_package_lazy_attrs.py +++ b/tests/unit/scenario/test_package_lazy_attrs.py @@ -3,19 +3,25 @@ """Tests for the lazy ``__getattr__`` hooks on scenario subpackages.""" -from unittest.mock import MagicMock +from collections.abc import Callable, Iterator +from importlib import import_module +from unittest.mock import MagicMock, patch import pytest +from pyrit.executor.attack import PromptSendingAttack from pyrit.prompt_target import PromptTarget -from pyrit.registry import TargetRegistry +from pyrit.registry import Registry, TargetRegistry from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry +from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory from pyrit.scenario.core.scenario_technique import ScenarioTechnique from pyrit.scenario.scenarios._dynamic_techniques import reset_dynamic_technique_caches from pyrit.scenario.scenarios.airt.cyber import _build_cyber_technique from pyrit.scenario.scenarios.airt.leakage import _build_leakage_technique from pyrit.scenario.scenarios.airt.rapid_response import _build_rapid_response_technique from pyrit.scenario.scenarios.benchmark.adversarial import _build_benchmark_technique +from pyrit.score import TrueFalseScorer +from pyrit.setup.initialization import reset_setup_registries from pyrit.setup.initializers.techniques import build_technique_factories @@ -94,3 +100,97 @@ def test_unknown_attribute_raises(self) -> None: with pytest.raises(AttributeError, match="no attribute 'NotAThing'"): _ = benchmark.NotAThing # type: ignore[attr-defined] + + +@pytest.fixture +def isolated_setup_registries() -> Iterator[None]: + with patch.dict(Registry._singletons): + try: + yield + finally: + reset_dynamic_technique_caches() + + +@pytest.mark.usefixtures("patch_central_database", "isolated_setup_registries") +class TestDynamicTechniqueReset: + @pytest.mark.parametrize("reset_caches", [reset_dynamic_technique_caches, reset_setup_registries]) + @pytest.mark.parametrize( + ("package_name", "scenario_name", "technique_name"), + [ + ("garak", "Doctor", "DoctorTechnique"), + ("airt", "Jailbreak", "JailbreakTechnique"), + ], + ) + def test_local_technique_identity_survives_reset( + self, + *, + package_name: str, + scenario_name: str, + technique_name: str, + reset_caches: Callable[[], None], + ) -> None: + packages = [ + import_module(f"pyrit.scenario.{package_name}"), + import_module(f"pyrit.scenario.scenarios.{package_name}"), + import_module("pyrit.scenario.scenarios._dynamic_techniques"), + ] + technique_class = getattr(packages[0], technique_name) + scenario_class = getattr(packages[0], scenario_name) + scorer = MagicMock(spec=TrueFalseScorer) + before = scenario_class(objective_scorer=scorer) + + for _ in range(2): + reset_caches() + + assert all(getattr(package, technique_name) is technique_class for package in packages) + after = scenario_class(objective_scorer=scorer) + assert after._technique_class is before._technique_class is technique_class + assert after._default_technique is before._default_technique is technique_class.default() + assert ( + after._technique_class.resolve([technique_class["ALL"]], default=after._default_technique) + == technique_class.get_all_techniques() + ) + + @pytest.mark.parametrize( + ("package_name", "technique_name"), + [ + ("airt", "CyberTechnique"), + ("airt", "LeakageTechnique"), + ("airt", "MultilingualTechnique"), + ("airt", "RapidResponseTechnique"), + ("benchmark", "AdversarialBenchmarkTechnique"), + ], + ) + def test_registry_technique_cache_and_exports_are_discarded( + self, *, package_name: str, technique_name: str + ) -> None: + package = import_module(f"pyrit.scenario.scenarios.{package_name}") + dynamic = import_module("pyrit.scenario.scenarios._dynamic_techniques") + reset_dynamic_technique_caches() + before = getattr(dynamic, technique_name) + assert getattr(package, technique_name) is before + + reset_dynamic_technique_caches() + + assert technique_name not in vars(package) + after = getattr(dynamic, technique_name) + assert after is not before + assert getattr(package, technique_name) is after + + def test_setup_reset_rebuilds_from_replacement_catalog(self) -> None: + import pyrit.scenario.scenarios.airt as airt + + reset_dynamic_technique_caches() + before = airt.RapidResponseTechnique + before_techniques = before.get_all_techniques() + reset_setup_registries() + factory = AttackTechniqueFactory(name="replacement", attack_class=PromptSendingAttack) + AttackTechniqueRegistry.get_registry_singleton().register_from_factories([factory]) + + after = airt.RapidResponseTechnique + + assert after is not before + assert [technique.value for technique in after.get_all_techniques()] == ["replacement"] + assert before.get_all_techniques() == before_techniques + with pytest.raises(ValueError, match="unsupported techniques"): + after.resolve([before["ALL"]], default=after.default()) diff --git a/tests/unit/setup/test_reinitialization.py b/tests/unit/setup/test_reinitialization.py index 22be363284..40530580cf 100644 --- a/tests/unit/setup/test_reinitialization.py +++ b/tests/unit/setup/test_reinitialization.py @@ -5,16 +5,47 @@ import os import uuid +from collections.abc import Generator, Iterator +from contextlib import contextmanager, nullcontext from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch import pytest +from pyrit.common.apply_defaults import get_global_default_values +from pyrit.common.random_context import configure_random_seed, get_configured_random_seed from pyrit.memory import AzureSQLMemory, CentralMemory, SQLiteMemory from pyrit.registry import AttackRegistry, InitializerRegistry, Registry, TargetRegistry from pyrit.setup.configuration_loader import ConfigurationLoader from pyrit.setup.environment_loading import resolve_environment_async from pyrit.setup.initialization import reset_setup_registries, validate_reinitialization_memory +from pyrit.setup.initializers.targets import TARGET_CONFIGS + + +@contextmanager +def _isolated_reinitialization_state() -> Generator[None, None, None]: + defaults = get_global_default_values() + previous_seed = get_configured_random_seed() + with ( + patch.dict(os.environ), + patch.dict(Registry._singletons, {}, clear=True), + patch.dict(defaults._default_values, {}, clear=True), + ): + for config in TARGET_CONFIGS: + for name in (config.endpoint_var, config.key_var, config.model_var, config.underlying_model_var): + if name: + os.environ.pop(name, None) + try: + yield + finally: + reset_setup_registries() + configure_random_seed(seed=previous_seed) + + +@pytest.fixture(autouse=True) +def isolate_reinitialization_state() -> Iterator[None]: + with _isolated_reinitialization_state(): + yield @pytest.mark.usefixtures("patch_central_database") @@ -90,27 +121,86 @@ async def initialize_async(self): reset_setup_registries() -async def test_replacement_precedence_interpolation_empty_and_omission(tmp_path: Path) -> None: +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize("configured_target", [False, True]) +async def test_replacement_precedence_interpolation_empty_and_omission( + *, tmp_path: Path, sqlite_instance: SQLiteMemory, configured_target: bool +) -> None: base, other, local = tmp_path / ".env", tmp_path / "other.env", tmp_path / ".env.local" base.write_text("VALUE=new\nEMPTY=\nINTERPOLATED=${VALUE}/suffix\n", encoding="utf-8") other.write_text("VALUE=ignored\n", encoding="utf-8") local.write_text("VALUE=local\n", encoding="utf-8") with patch.dict(os.environ, {"VALUE": "old", "EMPTY": "old", "OMITTED": "keep"}): + if configured_target: + os.environ.update( + { + "OPENAI_CHAT_ENDPOINT": "http://127.0.0.1:1", + "OPENAI_CHAT_KEY": "synthetic-test-key", + "OPENAI_CHAT_MODEL": "gpt-4o", + } + ) values = await resolve_environment_async(env_files=[base, other, local], env_akv_ref=None, env_akv_strict=True) assert values == {"VALUE": "local", "EMPTY": "", "INTERPOLATED": "local/suffix"} assert os.environ["VALUE"] == "old" - with ( - patch("pyrit.setup.initialization.validate_reinitialization_memory", return_value=object()), - patch.object(CentralMemory, "set_memory_instance"), - ): - config = ConfigurationLoader(memory_db_type="in_memory", initialization_scripts=[], env_files=[]) - prepared = await config.preflight_reinitialization_async(environment_values=values) - await config.apply_prepared_reinitialization_async(prepared=prepared) + config = ConfigurationLoader(memory_db_type="in_memory", initialization_scripts=[], env_files=[]) + prepared = await config.preflight_reinitialization_async(environment_values=values) + await config.apply_prepared_reinitialization_async(prepared=prepared) + assert CentralMemory.get_memory_instance() is sqlite_instance + targets = TargetRegistry.get_registry_singleton().instances + if configured_target: + assert targets.get("openai_chat")._memory is sqlite_instance + else: + assert targets.get_names() == [] assert os.environ["VALUE"] == "local" assert os.environ["EMPTY"] == "" assert os.environ["OMITTED"] == "keep" +async def test_reinitialization_without_memory_does_not_mutate_state_async(sqlite_instance: SQLiteMemory) -> None: + assert CentralMemory.get_memory_instance() is sqlite_instance + config = ConfigurationLoader(memory_db_type="in_memory", initialization_scripts=[], env_files=[]) + prepared = await config.preflight_reinitialization_async(environment_values={"DO_NOT_APPLY": "changed"}) + registry = TargetRegistry.get_registry_singleton() + initializers = InitializerRegistry.get_registry_singleton() + environment = dict(os.environ) + + with patch.object(CentralMemory, "_memory_instance", None): + with pytest.raises(RuntimeError, match="Live reinitialization requires initialized memory"): + await config.apply_prepared_reinitialization_async(prepared=prepared) + + assert dict(os.environ) == environment + assert TargetRegistry.get_registry_singleton() is registry + assert InitializerRegistry.get_registry_singleton() is initializers + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize("fail", [False, True]) +def test_reinitialization_scope_restores_state_after_success_or_failure(fail: bool) -> None: + registry = TargetRegistry.get_registry_singleton() + memory = CentralMemory.get_memory_instance() + defaults = get_global_default_values() + defaults.set_default_value(class_type=ConfigurationLoader, parameter_name="operation", value="outside") + expected_defaults = defaults.all_defaults + environment = dict(os.environ) + previous_seed = get_configured_random_seed() + outcome = pytest.raises(RuntimeError, match="scope failure") if fail else nullcontext() + + with outcome: + with _isolated_reinitialization_state(): + os.environ["PYRIT_REINIT_TEST_CAPTURED"] = "inside" + assert TargetRegistry.get_registry_singleton() is not registry + defaults.set_default_value(class_type=ConfigurationLoader, parameter_name="operation", value="inside") + configure_random_seed(seed=7) + if fail: + raise RuntimeError("scope failure") + + assert TargetRegistry.get_registry_singleton() is registry + assert CentralMemory.get_memory_instance() is memory + assert defaults.all_defaults == expected_defaults + assert dict(os.environ) == environment + assert get_configured_random_seed() == previous_seed + + async def test_key_vault_selection_skips_default_base_and_overrides_process(tmp_path: Path) -> None: (tmp_path / ".env").write_text("IGNORED=base\n", encoding="utf-8") (tmp_path / ".env.local").write_text("VALUE=local\n", encoding="utf-8")