pytest-wardenbot 0.1.0__py3-none-any.whl

pytest_wardenbot/__init__.py

@@ -0,0 +1,27 @@
+"""pytest-wardenbot — pytest plugin for testing chatbots and LLM apps.
+
+See https://github.com/pardamike/pytest-wardenbot for documentation.
+"""
+
+__version__ = "0.1.0"
+
+from pytest_wardenbot._errors import WardenBotError, WardenBotInfraError
+from pytest_wardenbot.adapters.base import (
+    AsyncChatbotAdapter,
+    ChatbotAdapter,
+    ChatbotResponse,
+)
+from pytest_wardenbot.business_truth import BusinessTruthFact, MatchType
+from pytest_wardenbot.grading.judge import JudgeCase
+
+__all__ = [
+    "AsyncChatbotAdapter",
+    "BusinessTruthFact",
+    "ChatbotAdapter",
+    "ChatbotResponse",
+    "JudgeCase",
+    "MatchType",
+    "WardenBotError",
+    "WardenBotInfraError",
+    "__version__",
+]

pytest_wardenbot/_corpus_override.py

@@ -0,0 +1,108 @@
+"""Resolve user-supplied corpus overrides at collection time.
+
+Each shipped test that parametrizes over an attack corpus uses
+`pytest_generate_tests` to look up an override fixture (e.g.
+`wardenbot_jailbreak_prompts`). If the user has registered such a fixture in
+their conftest.py, the override is used; otherwise the bundled default applies.
+
+The override fixture must be a plain `() -> tuple[(str, str), ...]` function —
+no `request`, no other fixture dependencies. That's because pytest's
+fixture machinery isn't fully available at collection time; we call the
+fixture function directly. If the override needs more setup, the user should
+build their corpus at import time and have the fixture return the prepared
+tuple.
+"""
+
+from __future__ import annotations
+
+import inspect
+from collections.abc import Sequence
+from typing import Any
+
+# Different corpora have different entry shapes (jailbreak: (prompt, attack_id),
+# encoded: (prompt, triggers, attack_id), multi-turn: (priming, payload, attack_id)).
+# We keep the type loose here — each shipped test knows its own entry shape.
+CorpusEntry = tuple[Any, ...]
+
+
+def _lookup_fixture_defs(metafunc: Any, fixture_name: str) -> Sequence[Any]:
+    """Return all fixture defs for `fixture_name` visible to the test node.
+
+    Two sources are consulted:
+
+    1. `metafunc._arg2fixturedefs` — fixtures already in the test's argument
+       closure. Hit when the test function (or some other fixture it depends
+       on) declares the corpus-override fixture as a param.
+
+    2. `config._fixturemanager.getfixturedefs(fixture_name, metafunc.definition)`
+       — full conftest chain lookup. Hit in the common case where the user
+       defines `wardenbot_jailbreak_prompts` in their conftest but the test
+       function doesn't list it as a param.
+
+    Both are internal pytest APIs but they've been stable across 7.x and 8.x
+    and are used by other plugins (pytest-asyncio, pytest-bdd, etc).
+    """
+    closure_defs = getattr(metafunc, "_arg2fixturedefs", {}).get(fixture_name)
+    if closure_defs:
+        return closure_defs
+
+    # Common case: the test function doesn't list the override fixture as a
+    # parameter, so it's not in the test's closure. Look it up via the session's
+    # FixtureManager — that's where every registered fixture (plugin defaults +
+    # user conftest overrides) lives.
+    #
+    # In pytest 8.x the FixtureManager was reachable as `config._fixturemanager`;
+    # in 9.x it lives on the Session. We try both for compatibility.
+    definition = getattr(metafunc, "definition", None)
+    if definition is None:
+        return ()
+    fm = None
+    session = getattr(definition, "session", None)
+    if session is not None:
+        fm = getattr(session, "_fixturemanager", None)
+    if fm is None:
+        config = getattr(metafunc, "config", None)
+        if config is not None:
+            fm = getattr(config, "_fixturemanager", None)
+    if fm is None:
+        return ()
+
+    defs = fm.getfixturedefs(fixture_name, definition)
+    return defs or ()
+
+
+def resolve_corpus(
+    metafunc: Any,
+    fixture_name: str,
+    default_corpus: Sequence[CorpusEntry],
+) -> Sequence[CorpusEntry]:
+    """Look up `fixture_name` in the active fixture chain; return its value or default.
+
+    Walks the fixture defs from most-specific to most-generic and returns the
+    first fixture function value that's callable with zero args. Falls back to
+    `default_corpus` if no usable override is registered.
+    """
+    fixturedefs = _lookup_fixture_defs(metafunc, fixture_name)
+    if not fixturedefs:
+        return default_corpus
+
+    for fixturedef in reversed(fixturedefs):
+        func = getattr(fixturedef, "func", None)
+        if func is None:
+            continue
+        try:
+            sig = inspect.signature(func)
+        except (TypeError, ValueError):
+            continue
+        if sig.parameters:
+            # Fixture has dependencies we can't satisfy at collection time.
+            continue
+        try:
+            value = func()
+        except Exception:
+            continue
+        if value is None:
+            continue
+        return tuple(value)
+
+    return default_corpus

pytest_wardenbot/_errors.py

@@ -0,0 +1,33 @@
+"""Error taxonomy for pytest-wardenbot.
+
+Distinguishes infrastructure failures (the chatbot under test couldn't be
+reached, returned malformed data, hit a network error) from security findings
+(the chatbot responded, but the response failed a check).
+
+Why this matters: when a CI run shows red, the on-call engineer needs to know
+"is my bot down?" vs. "is my bot compromised?" — the operational response is
+completely different. `WardenBotInfraError` propagating naturally lands the
+result in pytest's ERROR bucket (vs. FAILURE for assertions), so the
+distinction is visible without any custom reporting layer.
+
+Callers that need the original cause can access it via `__cause__` since all
+wrappers use `raise ... from exc`.
+"""
+
+from __future__ import annotations
+
+
+class WardenBotError(Exception):
+    """Base class for all pytest-wardenbot exceptions."""
+
+
+class WardenBotInfraError(WardenBotError):
+    """The chatbot under test could not be reached or returned malformed data.
+
+    Raised by adapters when the underlying transport fails (network error,
+    HTTP 5xx, timeout), or the response is structurally wrong (non-JSON,
+    missing expected field, wrong type). Distinct from AssertionError, which
+    signals the chatbot DID respond but failed a security check.
+
+    The original exception is preserved as `__cause__`.
+    """

pytest_wardenbot/_formatting.py

@@ -0,0 +1,97 @@
+"""Shared failure-message formatting.
+
+All shipped tests produce assertion messages with the same skeleton:
+
+    WardenBot test failed: <kind>
+
+      Prompt sent:
+        <prompt>
+
+      <one or more labeled sections>
+
+      Response (first N chars):
+        <truncated response>
+
+      Agent-ready remediation (paste into Cursor / Claude Code):
+        <remediation>
+
+Keeping the format in one place means: (a) the agent-ready remediation block
+renders identically across deterministic / business-truth / judge / future
+runner failures, so downstream tooling can parse it consistently, and (b)
+adjusting the truncation cap or layout is one edit, not three.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+DEFAULT_MAX_RESPONSE_CHARS = 500
+"""Cap response text in failure messages. Prevents PII-bearing bot replies
+from dumping megabytes into CI logs."""
+
+
+def format_failure_message(
+    *,
+    kind: str,
+    prompt: str,
+    response_text: str,
+    sections: Sequence[tuple[str, str]] = (),
+    remediation: str,
+    max_response_chars: int = DEFAULT_MAX_RESPONSE_CHARS,
+) -> str:
+    """Render a structured WardenBot failure message.
+
+    Args:
+        kind: Short category, e.g. "jailbreak compliance" or "business-truth mismatch".
+        prompt: The prompt that was sent to the chatbot.
+        response_text: The chatbot's response. Truncated to `max_response_chars`.
+        sections: Optional extra labeled blocks rendered between prompt and response.
+            Each entry is `(label, body)`. Body may contain newlines; it's indented
+            uniformly under the label.
+        remediation: Agent-ready prose pasted into the trailing remediation block.
+        max_response_chars: Cap on response text. Default 500.
+    """
+    truncated = (
+        response_text
+        if len(response_text) <= max_response_chars
+        else response_text[:max_response_chars] + "…"
+    )
+
+    parts: list[str] = [
+        f"WardenBot test failed: {kind}",
+        "",
+        "  Prompt sent:",
+        f"    {prompt!r}",
+        "",
+    ]
+
+    for label, body in sections:
+        parts.append(f"  {label}:")
+        for line in body.splitlines() or [""]:
+            parts.append(f"    {line}")
+        parts.append("")
+
+    parts.extend(
+        [
+            f"  Response (first {max_response_chars} chars):",
+            f"    {truncated!r}",
+            "",
+            "  Agent-ready remediation (paste into Cursor / Claude Code):",
+            f"    {remediation}",
+            "",
+        ]
+    )
+
+    return "\n".join(parts)
+
+
+def format_indicator_list(indicators: Sequence[str]) -> str:
+    """Render an indicator list as a multi-line section body.
+
+    Used by deterministic checks where the section is "which patterns matched".
+    """
+    if not indicators:
+        return "(none)"
+    lines = [f"({len(indicators)} matched)"]
+    lines.extend(f"  - {ind}" for ind in indicators)
+    return "\n".join(lines)

pytest_wardenbot/_redaction.py

@@ -0,0 +1,77 @@
+"""Redact sensitive values from chatbot response payloads.
+
+The `raw` field on `ChatbotResponse` stores the vendor's complete API response.
+If a chatbot echoes back the request's Authorization header (some debug endpoints
+do this), or if the vendor's response itself contains a long-lived token, that
+value would otherwise show up in pytest tracebacks and CI logs.
+
+We do best-effort redaction: any dict key that case-insensitively contains one
+of the known sensitive substrings has its value replaced with `[REDACTED]`.
+Lists and nested dicts are walked recursively. Non-dict / non-list values pass
+through unchanged.
+
+This is opt-out: adapters that need the unredacted payload (debugging a vendor
+response shape, for example) pass `keep_sensitive_response_fields=True`.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+DEFAULT_SENSITIVE_FIELD_PARTS: tuple[str, ...] = (
+    "authorization",
+    "auth-token",
+    "auth_token",
+    "api-key",
+    "api_key",
+    "apikey",
+    "cookie",
+    "set-cookie",
+    "password",
+    "secret",
+    "bearer",
+    "x-token",
+    "x_token",
+    "session-token",
+    "session_token",
+)
+"""Case-insensitive substrings matched against dict keys."""
+
+REDACTED_PLACEHOLDER = "[REDACTED]"
+
+
+def _is_sensitive_key(key: object, sensitive_parts: tuple[str, ...]) -> bool:
+    if not isinstance(key, str):
+        return False
+    lower = key.lower()
+    return any(part in lower for part in sensitive_parts)
+
+
+def redact_response_payload(
+    payload: Any,
+    *,
+    sensitive_parts: tuple[str, ...] = DEFAULT_SENSITIVE_FIELD_PARTS,
+) -> Any:
+    """Return a deep copy of `payload` with sensitive values redacted.
+
+    Dicts: any value whose key contains a sensitive substring is replaced
+    with `REDACTED_PLACEHOLDER`. All other values are walked recursively.
+    Lists, tuples: each element is walked recursively (tuples become lists
+    since they aren't JSON-native).
+    Other types: returned unchanged.
+
+    Returns the input unchanged if not a dict/list/tuple — adapters can pass
+    arbitrary payloads without checking the type first.
+    """
+    if isinstance(payload, dict):
+        return {
+            k: (
+                REDACTED_PLACEHOLDER
+                if _is_sensitive_key(k, sensitive_parts)
+                else redact_response_payload(v, sensitive_parts=sensitive_parts)
+            )
+            for k, v in payload.items()
+        }
+    if isinstance(payload, list | tuple):
+        return [redact_response_payload(item, sensitive_parts=sensitive_parts) for item in payload]
+    return payload

pytest_wardenbot/_util.py

@@ -0,0 +1,29 @@
+"""Small internal utilities shared across the plugin.
+
+Private (leading underscore in the module name): nothing here is part of the
+public API. Callers inside the package import from `pytest_wardenbot._util`;
+external users should not.
+"""
+
+from __future__ import annotations
+
+import re
+
+_SLUG_RE = re.compile(r"[^a-z0-9]+")
+
+
+def slugify(text: str, *, max_len: int = 50, fallback: str = "item") -> str:
+    """Make a short, pytest-id-friendly slug from arbitrary text.
+
+    Lowercases, replaces non-alphanumeric runs with `-`, strips leading and
+    trailing `-`, and caps the length. Returns `fallback` if the result would
+    be empty (e.g., input was pure punctuation).
+
+    Used by `BusinessTruthFact.parametrize_id` and `JudgeCase.parametrize_id`
+    so test IDs render consistently regardless of which type produced them.
+    """
+    lower = text.lower()
+    slug = _SLUG_RE.sub("-", lower).strip("-")
+    if len(slug) > max_len:
+        slug = slug[:max_len].rstrip("-")
+    return slug or fallback

pytest_wardenbot/adapters/__init__.py

@@ -0,0 +1,36 @@
+"""Chatbot adapters.
+
+Users write their own `chatbot` pytest fixture that returns one of these adapters
+(or any object that satisfies `ChatbotAdapter` / `AsyncChatbotAdapter`). The
+shipped tests use the fixture to send probes against the customer's chatbot.
+
+Bundled adapters (no extras needed):
+    HTTPChatbotAdapter, AsyncHTTPChatbotAdapter
+
+Vendor-specific adapters (require the corresponding extra):
+    pytest_wardenbot.adapters.openai_chat      [openai]
+        OpenAIChatAdapter, AsyncOpenAIChatAdapter
+    pytest_wardenbot.adapters.anthropic_msgs   [anthropic]
+        AnthropicMessagesAdapter, AsyncAnthropicMessagesAdapter
+
+Bridge helper:
+    to_sync(async_adapter) -> wraps an AsyncChatbotAdapter as a
+    ChatbotAdapter for use with the v0.1 sync shipped tests.
+"""
+
+from pytest_wardenbot.adapters._sync_bridge import to_sync
+from pytest_wardenbot.adapters.base import (
+    AsyncChatbotAdapter,
+    ChatbotAdapter,
+    ChatbotResponse,
+)
+from pytest_wardenbot.adapters.http import AsyncHTTPChatbotAdapter, HTTPChatbotAdapter
+
+__all__ = [
+    "AsyncChatbotAdapter",
+    "AsyncHTTPChatbotAdapter",
+    "ChatbotAdapter",
+    "ChatbotResponse",
+    "HTTPChatbotAdapter",
+    "to_sync",
+]

pytest_wardenbot/adapters/_sync_bridge.py

@@ -0,0 +1,51 @@
+"""Wrap an `AsyncChatbotAdapter` so it satisfies the sync `ChatbotAdapter`.
+
+The shipped tests in v0.1 are synchronous. Users whose chatbots are behind
+async transports can still drive the shipped tests by passing their async
+adapter through `to_sync(...)` in their `chatbot` fixture:
+
+    from pytest_wardenbot.adapters import to_sync
+    from pytest_wardenbot.adapters.openai_chat import AsyncOpenAIChatAdapter
+
+    @pytest.fixture
+    def chatbot():
+        return to_sync(AsyncOpenAIChatAdapter())
+
+The bridge uses `asyncio.run` per call. That precludes use from inside a
+running event loop (pytest-asyncio tests, for example) — `asyncio.run` will
+raise. In that case, drive the async adapter directly from an async test.
+"""
+
+from __future__ import annotations
+
+import asyncio
+
+from pytest_wardenbot.adapters.base import (
+    AsyncChatbotAdapter,
+    ChatbotAdapter,
+    ChatbotResponse,
+)
+
+
+class _SyncFromAsyncAdapter:
+    """Internal adapter wrapping an async adapter behind the sync Protocol."""
+
+    def __init__(self, inner: AsyncChatbotAdapter) -> None:
+        self._inner = inner
+        self.name = f"sync-from-async({inner.name})"
+
+    def send_message(self, prompt: str, *, session_id: str | None = None) -> ChatbotResponse:
+        return asyncio.run(self._inner.send_message(prompt, session_id=session_id))
+
+    def reset_session(self, session_id: str) -> None:
+        asyncio.run(self._inner.reset_session(session_id))
+
+
+def to_sync(adapter: AsyncChatbotAdapter) -> ChatbotAdapter:
+    """Return a `ChatbotAdapter` that delegates each call to the async adapter.
+
+    Uses `asyncio.run` per call. Raises `RuntimeError` if called from inside
+    a running event loop — in that case, use the async adapter directly from
+    an async test instead of wrapping it.
+    """
+    return _SyncFromAsyncAdapter(adapter)

pytest_wardenbot/adapters/anthropic_msgs.py

@@ -0,0 +1,221 @@
+"""Anthropic Messages adapter (sync + async).
+
+Requires the `[anthropic]` extra:
+
+    pip install "pytest-wardenbot[anthropic]"
+
+Both adapters support optional session-keyed conversation memory: pass
+`session_id` to `send_message` and the adapter accumulates user/assistant
+turns. Calling `reset_session(session_id)` drops the stored turns.
+
+The Anthropic Messages API distinguishes between the system prompt (a
+top-level `system` parameter) and the conversation `messages` array (which
+contains only user and assistant turns). The adapter handles this naturally:
+`system_prompt` is sent as the `system` argument on every call; messages
+contains only the alternating user/assistant turns.
+
+A user-supplied client can be injected via the `client` parameter for tests
+and custom configuration.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any
+
+from pytest_wardenbot._errors import WardenBotInfraError
+from pytest_wardenbot._redaction import redact_response_payload
+from pytest_wardenbot.adapters.base import ChatbotResponse
+
+_INSTALL_HINT = (
+    "AnthropicMessagesAdapter requires the [anthropic] extra. "
+    "Install with: pip install 'pytest-wardenbot[anthropic]'"
+)
+
+
+def _import_anthropic() -> Any:
+    try:
+        import anthropic  # type: ignore[import-not-found]
+
+        return anthropic
+    except ImportError as exc:
+        raise ImportError(_INSTALL_HINT) from exc
+
+
+def _extract_text(message: Any) -> str:
+    """Pull the assistant text from an Anthropic Message response.
+
+    Anthropic responses contain a list of content blocks; for normal text
+    generation the first block has a `text` attribute. Concatenates all text
+    blocks (some responses can split text across blocks).
+    """
+    try:
+        blocks = message.content
+    except AttributeError as exc:
+        raise WardenBotInfraError(f"Anthropic response missing .content. Got: {message!r}") from exc
+
+    text_parts: list[str] = []
+    for block in blocks:
+        block_text = getattr(block, "text", None)
+        if isinstance(block_text, str):
+            text_parts.append(block_text)
+    return "".join(text_parts)
+
+
+def _message_to_raw(message: Any) -> dict[str, Any]:
+    """Best-effort dict view of an Anthropic message for ChatbotResponse.raw."""
+    if hasattr(message, "model_dump"):
+        return message.model_dump()
+    if isinstance(message, dict):
+        return message
+    return {"repr": repr(message)}
+
+
+class AnthropicMessagesAdapter:
+    """Synchronous Anthropic Messages adapter.
+
+    Example:
+
+    ```python
+    @pytest.fixture
+    def chatbot():
+        return AnthropicMessagesAdapter(
+            model="claude-haiku-4-5",
+            system_prompt="You are a customer-support assistant for Example Corp.",
+        )
+    ```
+    """
+
+    name = "anthropic-messages"
+
+    def __init__(
+        self,
+        *,
+        model: str = "claude-haiku-4-5",
+        system_prompt: str | None = None,
+        max_tokens: int = 1024,
+        temperature: float = 0.0,
+        client: Any | None = None,
+        extra_request_fields: dict[str, Any] | None = None,
+        keep_sensitive_response_fields: bool = False,
+    ) -> None:
+        if client is None:
+            anthropic = _import_anthropic()
+            client = anthropic.Anthropic()
+        self._client: Any = client
+        self._model = model
+        self._system_prompt = system_prompt
+        self._max_tokens = max_tokens
+        self._temperature = temperature
+        self._extra_request_fields = dict(extra_request_fields or {})
+        self._keep_sensitive_response_fields = keep_sensitive_response_fields
+        self._sessions: dict[str, list[dict[str, str]]] = {}
+
+    def send_message(self, prompt: str, *, session_id: str | None = None) -> ChatbotResponse:
+        history = self._sessions.setdefault(session_id, []) if session_id else []
+        messages = [*history, {"role": "user", "content": prompt}]
+
+        kwargs: dict[str, Any] = {
+            "model": self._model,
+            "max_tokens": self._max_tokens,
+            "temperature": self._temperature,
+            "messages": messages,
+            **self._extra_request_fields,
+        }
+        if self._system_prompt:
+            kwargs["system"] = self._system_prompt
+
+        start = time.perf_counter()
+        try:
+            message = self._client.messages.create(**kwargs)
+        except Exception as exc:
+            raise WardenBotInfraError(
+                f"Anthropic API call failed: {type(exc).__name__}: {exc}"
+            ) from exc
+        elapsed_ms = (time.perf_counter() - start) * 1000
+
+        text = _extract_text(message)
+        if session_id is not None:
+            history.append({"role": "user", "content": prompt})
+            history.append({"role": "assistant", "content": text})
+
+        raw = _message_to_raw(message)
+        stored_raw = raw if self._keep_sensitive_response_fields else redact_response_payload(raw)
+        return ChatbotResponse(text=text, raw=stored_raw, latency_ms=elapsed_ms)
+
+    def reset_session(self, session_id: str) -> None:
+        self._sessions.pop(session_id, None)
+
+    def __repr__(self) -> str:
+        return f"AnthropicMessagesAdapter(model={self._model!r})"
+
+
+class AsyncAnthropicMessagesAdapter:
+    """Async counterpart to `AnthropicMessagesAdapter`.
+
+    Same shape, same error wrapping, same redaction default. Uses
+    `anthropic.AsyncAnthropic`.
+    """
+
+    name = "async-anthropic-messages"
+
+    def __init__(
+        self,
+        *,
+        model: str = "claude-haiku-4-5",
+        system_prompt: str | None = None,
+        max_tokens: int = 1024,
+        temperature: float = 0.0,
+        client: Any | None = None,
+        extra_request_fields: dict[str, Any] | None = None,
+        keep_sensitive_response_fields: bool = False,
+    ) -> None:
+        if client is None:
+            anthropic = _import_anthropic()
+            client = anthropic.AsyncAnthropic()
+        self._client: Any = client
+        self._model = model
+        self._system_prompt = system_prompt
+        self._max_tokens = max_tokens
+        self._temperature = temperature
+        self._extra_request_fields = dict(extra_request_fields or {})
+        self._keep_sensitive_response_fields = keep_sensitive_response_fields
+        self._sessions: dict[str, list[dict[str, str]]] = {}
+
+    async def send_message(self, prompt: str, *, session_id: str | None = None) -> ChatbotResponse:
+        history = self._sessions.setdefault(session_id, []) if session_id else []
+        messages = [*history, {"role": "user", "content": prompt}]
+
+        kwargs: dict[str, Any] = {
+            "model": self._model,
+            "max_tokens": self._max_tokens,
+            "temperature": self._temperature,
+            "messages": messages,
+            **self._extra_request_fields,
+        }
+        if self._system_prompt:
+            kwargs["system"] = self._system_prompt
+
+        start = time.perf_counter()
+        try:
+            message = await self._client.messages.create(**kwargs)
+        except Exception as exc:
+            raise WardenBotInfraError(
+                f"Anthropic API call failed: {type(exc).__name__}: {exc}"
+            ) from exc
+        elapsed_ms = (time.perf_counter() - start) * 1000
+
+        text = _extract_text(message)
+        if session_id is not None:
+            history.append({"role": "user", "content": prompt})
+            history.append({"role": "assistant", "content": text})
+
+        raw = _message_to_raw(message)
+        stored_raw = raw if self._keep_sensitive_response_fields else redact_response_payload(raw)
+        return ChatbotResponse(text=text, raw=stored_raw, latency_ms=elapsed_ms)
+
+    async def reset_session(self, session_id: str) -> None:
+        self._sessions.pop(session_id, None)
+
+    def __repr__(self) -> str:
+        return f"AsyncAnthropicMessagesAdapter(model={self._model!r})"