TEST Isolate SDK unit tests and move LiteLLM wire checks to integration

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.
 

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`.

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.

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)

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 == []

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
 
 

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="  ")]