pytest-resilience-agent 0.2.0__py3-none-any.whl

pytest_resilience_agent/__init__.py

@@ -0,0 +1,8 @@
+"""pytest-resilience-agent: pytest plugin for LLM resilience testing.
+
+Generates and runs tests that prove your LLM application keeps working when
+infrastructure breaks: gateway timeouts, model brownouts, MCP server errors,
+rate limits, partial outages.
+"""
+
+__version__ = "0.2.0"

pytest_resilience_agent/chaos.py

@@ -0,0 +1,174 @@
+"""Chaos scenario controller - real implementation.
+
+The controller owns a single ``respx.MockRouter`` and one scenario object
+per requested name. On enter, it starts the mock and applies every
+scenario. On exit, it reverts each scenario (so cleanup is deterministic)
+and stops the mock.
+
+OpenTelemetry: every scenario activation / deactivation / call interception
+emits a span on the ``pytest-resilience-agent`` tracer so downstream
+observability tooling sees the chaos events alongside real spans.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import pytest
+import respx
+from opentelemetry import trace
+
+from pytest_resilience_agent.scenarios import (
+    Scenario,
+    build_scenario,
+    registered_scenarios,
+)
+
+_TRACER = trace.get_tracer("pytest-resilience-agent")
+
+
+@dataclass
+class ChaosEvent:
+    """One chaos event captured for the test report."""
+
+    scenario: str
+    detail: str
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+class ChaosController:
+    """Apply a set of chaos scenarios against a target HTTP host for one test.
+
+    Default target URL is ``https://gateway.local/v1/chat/completions``.
+    Override with ``target_gateway_url`` / ``target_lark_url`` when the test
+    fixtures know real URLs.
+    """
+
+    DEFAULT_GATEWAY_URL = "https://gateway.local/v1/chat/completions"
+    DEFAULT_LARK_URL = "https://lark.local"
+
+    def __init__(
+        self,
+        scenarios: list[str] | None = None,
+        target_gateway_url: str | None = None,
+        target_lark_url: str | None = None,
+        turns: list[list[str]] | None = None,
+    ) -> None:
+        if scenarios and turns is not None:
+            raise pytest.UsageError(
+                "resilience marker accepts either scenarios= or turns=, not both"
+            )
+        if turns is not None:
+            if not turns:
+                raise pytest.UsageError("turns= must list at least one turn")
+            known = set(registered_scenarios())
+            unknown = [n for turn in turns for n in turn if n not in known]
+            if unknown:
+                raise pytest.UsageError(
+                    f"unknown chaos scenario(s) in turns=: {sorted(set(unknown))}. "
+                    f"Registered: {sorted(known)}"
+                )
+        self.scenario_names = list(scenarios) if scenarios else []
+        self.turns = turns
+        self.target_gateway_url = target_gateway_url or self.DEFAULT_GATEWAY_URL
+        self.target_lark_url = target_lark_url or self.DEFAULT_LARK_URL
+        self.events: list[ChaosEvent] = []
+        self._mock = respx.mock(assert_all_called=False, assert_all_mocked=False)
+        self._scenarios: list[Scenario] = []
+        self._turn_index = 0
+
+    @property
+    def current_turn(self) -> int:
+        """Zero-based index of the active conversation turn (0 in single-window mode)."""
+        return self._turn_index
+
+    def _apply_scenarios(self, names: list[str]) -> None:
+        """Build, apply and record one set of scenarios against the live mock."""
+        for name in names:
+            target = self._target_for(name)
+            scenario = build_scenario(name, self._mock, target)
+            result = scenario.apply()
+            self._scenarios.append(scenario)
+            self.events.append(
+                ChaosEvent(
+                    scenario=result.scenario,
+                    detail=result.detail,
+                    metadata=result.metadata,
+                )
+            )
+            with _TRACER.start_as_current_span(f"chaos.apply.{name}") as span:
+                for k, v in result.metadata.items():
+                    span.set_attribute(f"chaos.{k}", v)
+
+    def _apply_turn(self, index: int) -> None:
+        """Activate the scenarios for one conversation turn and mark the boundary."""
+        assert self.turns is not None
+        self._turn_index = index
+        names = list(self.turns[index])
+        with _TRACER.start_as_current_span(f"chaos.turn.{index}") as span:
+            span.set_attribute("chaos.turn.scenarios", names)
+            # Apply inside the turn span so the chaos.apply.* spans nest under it.
+            self._apply_scenarios(names)
+
+    def _revert_current(self) -> None:
+        """Revert (record stats for) every scenario active in the current turn."""
+        for scenario in reversed(self._scenarios):
+            result = scenario.revert()
+            self.events.append(
+                ChaosEvent(
+                    scenario=result.scenario,
+                    detail=result.detail,
+                    metadata=result.metadata,
+                )
+            )
+            with _TRACER.start_as_current_span(f"chaos.revert.{result.scenario}") as span:
+                for k, v in result.metadata.items():
+                    span.set_attribute(f"chaos.{k}", v)
+        self._scenarios.clear()
+
+    def next_turn(self) -> None:
+        """Advance to the next conversation turn: revert this turn's chaos, drop the
+        mock's routes and call history, then apply the next turn's scenarios.
+
+        Counters reset because each turn builds brand-new Scenario instances (each
+        starts at zero calls). Clearing the mock removes the previous turn's routes
+        and the accumulated call records, so a long conversation does not grow them
+        unbounded."""
+        if self.turns is None:
+            raise pytest.UsageError("next_turn() requires turns= on the resilience marker")
+        if self._turn_index + 1 >= len(self.turns):
+            raise pytest.UsageError(
+                f"next_turn() advanced past the last defined turn "
+                f"(have {len(self.turns)} turns, already at turn {self._turn_index})"
+            )
+        self._revert_current()
+        self._mock.reset()  # drop accumulated call records
+        self._mock.clear()  # drop the previous turn's routes
+        self._apply_turn(self._turn_index + 1)
+
+    def enter(self) -> None:
+        """Start respx and install scenarios for turn 0 (or the single window)."""
+        self._mock.start()
+        if self.turns is not None:
+            self._apply_turn(0)
+        else:
+            self._apply_scenarios(self.scenario_names)
+
+    def exit(self) -> None:
+        """Revert every scenario in LIFO order and stop respx."""
+        self._revert_current()
+        self._mock.stop()
+
+    def record(self, scenario: str, detail: str, **metadata: Any) -> None:
+        """Append a custom chaos event from inside user code."""
+        self.events.append(ChaosEvent(scenario=scenario, detail=detail, metadata=metadata))
+
+    def _target_for(self, scenario_name: str) -> str:
+        """Pick gateway URL or Lark URL depending on which layer the scenario hits."""
+        if scenario_name == "mcp_error":
+            return self.target_lark_url
+        return self.target_gateway_url
+
+
+__all__ = ["ChaosController", "ChaosEvent", "registered_scenarios"]

pytest_resilience_agent/cli.py

@@ -0,0 +1,199 @@
+"""Command-line tool entry point.
+
+Subcommands:
+
+- ``discover``     list failing tests via Lark MCP (or local mock)
+- ``generate``     synthesize resilience test files from those failures
+- ``run``          execute the generated tests via pytest
+- ``report``       push resolution status back to Lark
+- ``scenarios``    print the built-in chaos scenario registry
+
+Install with ``pip install -e .`` and call ``pytest-resilience-agent``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+from rich.console import Console
+from rich.table import Table
+
+from pytest_resilience_agent.generator import generate_test
+from pytest_resilience_agent.lark_mcp import LarkMCPClient
+from pytest_resilience_agent.scenarios import registered_scenarios
+
+_DEFAULT_LARK_URL = "http://localhost:8801"
+_DEFAULT_PROJECT = "demo"
+_DEFAULT_GENERATED_DIR = Path("generated_resilience_tests")
+
+
+def _console() -> Console:
+    """Return a Console that handles UTF-8 cleanly on Windows cp1252 terminals."""
+    return Console(force_terminal=True)
+
+
+def _lark_url(args: argparse.Namespace) -> str:
+    return args.lark_url or os.environ.get("LARK_MCP_URL") or _DEFAULT_LARK_URL
+
+
+# ---------------------------------------------------------------------------
+# discover
+# ---------------------------------------------------------------------------
+
+
+def cmd_discover(args: argparse.Namespace) -> int:
+    """List failing tests from the Lark MCP server."""
+    console = _console()
+    client = LarkMCPClient(base_url=_lark_url(args))
+    try:
+        tests = client.list_failing_tests(args.project)
+    finally:
+        client.close()
+    table = Table(title=f"Failing tests in {args.project!r}")
+    table.add_column("Name", style="cyan")
+    table.add_column("Last failure", style="red")
+    table.add_column("Path", style="dim")
+    for t in tests:
+        table.add_row(t.name, t.last_failure or "", t.path)
+    console.print(table)
+    if not tests:
+        console.print("[yellow]No failing tests reported.[/]")
+    return 0
+
+
+# ---------------------------------------------------------------------------
+# generate
+# ---------------------------------------------------------------------------
+
+
+def cmd_generate(args: argparse.Namespace) -> int:
+    """Pull failing tests from Lark and write resilience test files."""
+    console = _console()
+    client = LarkMCPClient(base_url=_lark_url(args))
+    try:
+        tests = client.list_failing_tests(args.project)
+    finally:
+        client.close()
+    if not tests:
+        console.print("[yellow]No failing tests, nothing to generate.[/]")
+        return 0
+    out_dir = Path(args.out)
+    table = Table(title=f"Generated {len(tests)} resilience test(s)")
+    table.add_column("Test name", style="cyan")
+    table.add_column("Scenarios", style="green")
+    table.add_column("File", style="dim")
+    for t in tests:
+        generated = generate_test(t.name, t.last_failure or "", out_dir)
+        table.add_row(t.name, ", ".join(generated.scenarios), str(generated.file_path))
+    console.print(table)
+    console.print(f"[green]Wrote {len(tests)} file(s) under {out_dir}[/]")
+    return 0
+
+
+# ---------------------------------------------------------------------------
+# run
+# ---------------------------------------------------------------------------
+
+
+def cmd_run(args: argparse.Namespace) -> int:
+    """Run pytest against the generated tests directory."""
+    target = Path(args.path)
+    if not target.exists():
+        _console().print(f"[red]No such path: {target}[/]")
+        return 2
+    cmd = [sys.executable, "-X", "utf8", "-m", "pytest", "-v", str(target)]
+    if args.record:
+        cmd.append(f"--resilience-record={args.record}")
+    result = subprocess.run(cmd)
+    return result.returncode
+
+
+# ---------------------------------------------------------------------------
+# report
+# ---------------------------------------------------------------------------
+
+
+def cmd_report(args: argparse.Namespace) -> int:
+    """Push a resolution back to Lark."""
+    client = LarkMCPClient(base_url=_lark_url(args))
+    try:
+        client.report_resolved(
+            project=args.project,
+            test_name=args.test_name,
+            pytest_path=args.pytest_path,
+        )
+    finally:
+        client.close()
+    _console().print(f"[green]Reported resolution for {args.test_name}[/]")
+    return 0
+
+
+# ---------------------------------------------------------------------------
+# scenarios
+# ---------------------------------------------------------------------------
+
+
+def cmd_scenarios(args: argparse.Namespace) -> int:
+    """Print the registered chaos scenario names."""
+    console = _console()
+    table = Table(title="Registered chaos scenarios")
+    table.add_column("Name", style="cyan")
+    for name in registered_scenarios():
+        table.add_row(name)
+    console.print(table)
+    return 0
+
+
+# ---------------------------------------------------------------------------
+# parser
+# ---------------------------------------------------------------------------
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """Construct the argparse tree."""
+    parser = argparse.ArgumentParser(prog="pytest-resilience-agent")
+    parser.add_argument(
+        "--lark-url",
+        default=None,
+        help="Lark MCP base URL (env LARK_MCP_URL, default http://localhost:8801)",
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    p_disc = sub.add_parser("discover", help="list failing tests via Lark")
+    p_disc.add_argument("--project", default=_DEFAULT_PROJECT)
+    p_disc.set_defaults(func=cmd_discover)
+
+    p_gen = sub.add_parser("generate", help="generate resilience tests from Lark failures")
+    p_gen.add_argument("--project", default=_DEFAULT_PROJECT)
+    p_gen.add_argument("--out", default=str(_DEFAULT_GENERATED_DIR))
+    p_gen.set_defaults(func=cmd_generate)
+
+    p_run = sub.add_parser("run", help="pytest the generated tests")
+    p_run.add_argument("--path", default=str(_DEFAULT_GENERATED_DIR))
+    p_run.add_argument("--record", default=None, help="JSON timeline path")
+    p_run.set_defaults(func=cmd_run)
+
+    p_rep = sub.add_parser("report", help="push a resolution back to Lark")
+    p_rep.add_argument("--project", default=_DEFAULT_PROJECT)
+    p_rep.add_argument("--test-name", required=True)
+    p_rep.add_argument("--pytest-path", required=True)
+    p_rep.set_defaults(func=cmd_report)
+
+    p_sce = sub.add_parser("scenarios", help="list registered chaos scenarios")
+    p_sce.set_defaults(func=cmd_scenarios)
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    return int(args.func(args) or 0)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

pytest_resilience_agent/gateway.py

@@ -0,0 +1,75 @@
+"""Thin client around TrueFoundry AI Gateway.
+
+The gateway is an OpenAI-compatible proxy that handles fallbacks, retries,
+and multi-model routing on its side. We just send chat completions through
+it and let the gateway config decide what happens when upstream models error.
+
+Docs: https://www.truefoundry.com/docs/ai-gateway/intro-to-llm-gateway
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import httpx
+
+
+@dataclass
+class ChatReply:
+    """Minimal reply object returned by ``AIGatewayClient.chat``."""
+
+    content: str
+    model: str
+    raw: dict[str, Any]
+
+
+class AIGatewayClient:
+    """Send chat completions through a TrueFoundry-style AI Gateway.
+
+    The gateway is OpenAI-compatible, so callers can keep the OpenAI SDK
+    interface in their own code and swap the base URL. We keep a thin
+    httpx client here to make timeout and error injection easy for tests.
+    """
+
+    def __init__(
+        self,
+        base_url: str,
+        api_key: str | None = None,
+        timeout: float = 30.0,
+    ) -> None:
+        self.base_url = base_url.rstrip("/")
+        self._timeout = timeout
+        headers: dict[str, str] = {
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+            "User-Agent": "pytest-resilience-agent/0.1.0 (httpx)",
+        }
+        if api_key:
+            headers["Authorization"] = f"Bearer {api_key}"
+        self._client = httpx.Client(headers=headers, timeout=timeout)
+
+    def chat(
+        self,
+        messages: list[dict[str, str]],
+        model: str = "gateway-default",
+        **kwargs: Any,
+    ) -> ChatReply:
+        """Send a chat completion request through the gateway.
+
+        The gateway's own config decides fallback chain and retries.
+        If the gateway itself is unreachable, raises ``httpx.HTTPError``.
+        """
+        payload = {"model": model, "messages": messages, **kwargs}
+        response = self._client.post(f"{self.base_url}/chat/completions", json=payload)
+        response.raise_for_status()
+        data = response.json()
+        choices = data.get("choices") or []
+        content = ""
+        if choices:
+            content = choices[0].get("message", {}).get("content", "") or ""
+        return ChatReply(content=content, model=data.get("model", model), raw=data)
+
+    def close(self) -> None:
+        """Close the underlying httpx client."""
+        self._client.close()

pytest_resilience_agent/generator.py

@@ -0,0 +1,118 @@
+"""Generate resilience test files from observed failures.
+
+Mirrors the Lark sponsor challenge example ("coding agent setup that
+listens for failing tests and creates PRs to fix them"). For each
+failing test reported by Lark, the generator picks the relevant chaos
+scenario(s) and emits a runnable pytest file that reproduces the
+failure under controlled conditions.
+
+The mapping is deterministic and rule-based on the failure string.
+That keeps generation auditable: the engineer can read the rule that
+matched and tweak it. A future iteration can swap in an LLM-based
+classifier behind the same Generator interface.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from pathlib import Path
+
+# Order matters: first match wins.
+_RULES: list[tuple[re.Pattern[str], list[str]]] = [
+    (re.compile(r"\b429\b|too many requests", re.I), ["rate_limit"]),
+    (re.compile(r"\b502\b|bad gateway", re.I), ["llm_5xx"]),
+    (re.compile(r"\b503\b|service unavailable", re.I), ["partial_outage"]),
+    (re.compile(r"\b504\b|timeout", re.I), ["llm_timeout"]),
+    (re.compile(r"\b402\b|quota|cost", re.I), ["cost_exceeded"]),
+    (re.compile(r"connect.*(error|refused|reset)", re.I), ["network_blip"]),
+    (re.compile(r"empty|truncated|stream.*drop", re.I), ["stream_stall"]),
+    (re.compile(r"mcp|tool.*error|jsonrpc", re.I), ["mcp_error"]),
+    (re.compile(r"wrong.*model|model.*mismatch", re.I), ["wrong_model_returned"]),
+]
+
+
+@dataclass
+class GeneratedTest:
+    """Result of generating one test file from a Lark failure."""
+
+    test_name: str
+    scenarios: list[str]
+    file_path: Path
+    source: str
+
+
+def pick_scenarios(failure_text: str) -> list[str]:
+    """Return the scenario names that best match a free-text failure.
+
+    Falls back to ``llm_5xx`` if nothing matches; that is the most generic
+    transient-failure scenario and a sane default.
+    """
+    matched: list[str] = []
+    for pattern, scenarios in _RULES:
+        if pattern.search(failure_text):
+            matched.extend(scenarios)
+    return matched or ["llm_5xx"]
+
+
+_TEMPLATE = '''"""Auto-generated resilience test for {test_name!r}.
+
+Generated by pytest-resilience-agent from a Lark-reported failure.
+
+Original failure text:
+    {failure_text}
+
+Chosen scenarios: {scenarios}
+"""
+
+from __future__ import annotations
+
+import httpx
+import pytest
+
+
+@pytest.mark.resilience(scenarios={scenarios!r})
+def {function_name}(chaos):
+    """Reproduce the original failure under controlled chaos.
+
+    The chaos fixture installs the named scenarios at the HTTP layer.
+    Replace the body below with the call from your real agent code; the
+    boilerplate POST below just exercises the gateway to prove the
+    scenarios are active.
+    """
+    with httpx.Client() as client:
+        response = client.post(chaos.target_gateway_url, json={{"q": 1}})
+        # When the scenarios are correct, the response status / shape
+        # mirrors the original failure. Add an assertion on the contract
+        # your real agent should hold:
+        # assert agent.handle_response(response).fallback_used
+        assert response is not None
+'''
+
+
+def generate_test(
+    test_name: str,
+    failure_text: str,
+    out_dir: Path,
+) -> GeneratedTest:
+    """Write a resilience test file for one failure. Return the artefact."""
+    scenarios = pick_scenarios(failure_text)
+    safe_name = re.sub(r"[^a-zA-Z0-9_]+", "_", test_name).strip("_") or "auto_test"
+    if not safe_name.startswith("test_"):
+        safe_name = f"test_{safe_name}"
+    file_path = out_dir / f"{safe_name}_resilience.py"
+    function_name = safe_name + "_resilience"
+    source = _TEMPLATE.format(
+        test_name=test_name,
+        failure_text=failure_text,
+        scenarios=scenarios,
+        function_name=function_name,
+    )
+    out_dir.mkdir(parents=True, exist_ok=True)
+    file_path.write_text(source, encoding="utf-8")
+    return GeneratedTest(
+        test_name=test_name,
+        scenarios=scenarios,
+        file_path=file_path,
+        source=source,
+    )