mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
08fa25042c
The shared proxy wrapper in tests/e2e/e2e_gateway.py was misnamed: Gateway is not a gateway server, it is the client every suite uses to talk to the proxy (keys, models, chat/embed/ocr, spend read-backs, poll helpers). Rename the module to proxy_client.py and the class to ProxyClient, with build_gateway becoming build_proxy_client and the GatewayProvider protocol becoming ProxyClientProvider. The .gateway attribute suites held is now .proxy. Only identifiers changed; prose and string literals that use the word gateway for the proxy-server concept were left alone. Each suite previously built its own instance through a per-suite build_client() that called build_gateway() inside, duplicating the proxy wiring across suites. There is now one session-scoped proxy fixture in tests/e2e/conftest.py; every suite's client fixture depends on it and injects it, so the wiring lives in one place. claude_code keeps building its own client directly since it has its own harness and does not use the shared fixtures. Behavior is unchanged: shared transport, data-plane/control-plane split routing, poll budget, typed request/response models, and resource cleanup all go through the same object.
120 lines
4.4 KiB
Python
120 lines
4.4 KiB
Python
"""Lifecycle contract and resource cleanup for stateful e2e tests.
|
|
|
|
Shared by every e2e suite under tests/e2e/. The proxy under test is
|
|
long-lived and never reset between tests, so anything a test creates (keys,
|
|
customers, teams, orgs, users, guardrails, budgets, ...) persists unless
|
|
explicitly deleted. Every check follows an init -> run -> teardown lifecycle;
|
|
teardown releases each resource init() created, even when run() raises.
|
|
|
|
In pytest terms (see conftest.py): the `resources` fixture's setup is init(),
|
|
the test body is run(), and the fixture's teardown is teardown().
|
|
"""
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Callable, List, Protocol, runtime_checkable
|
|
|
|
from proxy_client import ProxyClient
|
|
from models import KeyGenerateBody
|
|
|
|
|
|
@runtime_checkable
|
|
class E2ECase(Protocol):
|
|
"""A stateful e2e check run against a long-lived proxy.
|
|
|
|
init() acquires resources, run() exercises behaviour and asserts, teardown()
|
|
releases everything init() created. teardown() must run even if init() fails
|
|
partway or run() raises.
|
|
"""
|
|
|
|
def init(self) -> None: ...
|
|
|
|
def run(self) -> None: ...
|
|
|
|
def teardown(self) -> None: ...
|
|
|
|
|
|
def run_case(case: E2ECase) -> None:
|
|
"""Drive a case through its lifecycle: init -> run -> teardown.
|
|
|
|
teardown always runs - even when init() fails partway or run() raises (or
|
|
skips) - so resources the case already registered on the long-lived proxy are
|
|
released. init() is inside the try because cases register cleanups
|
|
progressively (e.g. create team, then user, then key), and a failure after
|
|
the first creation must still release what came before.
|
|
"""
|
|
try:
|
|
case.init()
|
|
case.run()
|
|
finally:
|
|
case.teardown()
|
|
|
|
|
|
@runtime_checkable
|
|
class ResourceClient(Protocol):
|
|
"""Proxy operations the convenience creators use. Resource types without a
|
|
creator here are handled generically via ResourceManager.defer(). The ProxyClient
|
|
satisfies this."""
|
|
|
|
def generate_key(self, body: KeyGenerateBody) -> str: ...
|
|
|
|
def delete_key(self, key: str) -> None: ...
|
|
|
|
def delete_customers(self, user_ids: List[str]) -> None: ...
|
|
|
|
|
|
@runtime_checkable
|
|
class ProxyClientProvider(Protocol):
|
|
"""Every suite's client exposes the shared ProxyClient, which the resources fixture
|
|
uses for cleanup. The client adds its own route methods on top."""
|
|
|
|
@property
|
|
def proxy(self) -> ProxyClient: ...
|
|
|
|
|
|
@dataclass
|
|
class ResourceManager:
|
|
"""Registry of teardown actions for resources a test creates on the stateful
|
|
proxy.
|
|
|
|
Not limited to any resource type: register a cleanup with ``defer()`` for a
|
|
key, customer, team, org, user, guardrail, budget, MCP server - anything with
|
|
a delete. The two most common resources have sugar (``key``, ``customer``);
|
|
everything else is ``resources.defer(lambda: client.delete_team(team_id))``.
|
|
|
|
Cleanups run LIFO (so a resource is removed before whatever it depends on) and
|
|
best-effort (one failing cleanup never blocks the rest).
|
|
"""
|
|
|
|
client: ResourceClient
|
|
_cleanups: List[Callable[[], None]] = field(
|
|
default_factory=list
|
|
) # mutable-ok: append-only teardown registry
|
|
|
|
def init(self) -> None:
|
|
"""No global setup needed today; present for lifecycle symmetry."""
|
|
return None
|
|
|
|
def defer(self, cleanup: Callable[[], None]) -> None:
|
|
"""Register a teardown action for any resource the test just created."""
|
|
self._cleanups.append(cleanup)
|
|
|
|
def key(self, models: list[str] | None = None, user_id: str | None = "e2e-test-user") -> str:
|
|
"""Create a virtual key; delete it on teardown. `models` restricts which
|
|
models the key may call (None/[] means all). `user_id` is required for
|
|
managed-batch ACL: the proxy stores created_by=user_id and checks it on
|
|
retrieve/cancel; None here means the 403 guard fires."""
|
|
key = self.client.generate_key(KeyGenerateBody(models=models or [], user_id=user_id))
|
|
self.defer(lambda: self.client.delete_key(key))
|
|
return key
|
|
|
|
def customer(self, customer_id: str) -> str:
|
|
"""Track an end-user id (from the `user` param); delete it on teardown."""
|
|
self.defer(lambda: self.client.delete_customers([customer_id]))
|
|
return customer_id
|
|
|
|
def teardown(self) -> None:
|
|
for cleanup in reversed(self._cleanups):
|
|
try:
|
|
cleanup()
|
|
except Exception:
|
|
pass # best-effort: a failed cleanup must not block the rest
|