mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
test(e2e): rename Gateway to ProxyClient and expose it as a session-scoped fixture (#33750)
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.
This commit is contained in:
Yassin Kortam
GitHub
parent
6d5f24fe0b
commit
08fa25042c
56 changed files with 441 additions and 423 deletions
|
|
@ -33,7 +33,7 @@ class TestPromptCompression:
|
|||
|
||||
def test_prompt_compression_accumulate_spend(self, key_id, user_id):
|
||||
for _ in range(10):
|
||||
response = self.resources.gateway.post("gemini-2.5-flash", key_id, user_id)
|
||||
response = self.resources.proxy.post("gemini-2.5-flash", key_id, user_id)
|
||||
compressed_value = ...
|
||||
assert response.cost == compressed_value # the cost was actually reduced
|
||||
```
|
||||
|
|
@ -48,9 +48,9 @@ The shape is layered so tests stay declarative
|
|||
|
||||
`transport.py` exposes a `Transport` Protocol with `post`, `get`, `delete`, `send`, `stream`, `probe`, plus `bearer(key)` and the `master` header. `HttpTransport` fulfils it, and `SplitTransport` routes each call by path to the data plane or the control plane so a split control-plane/data-plane deployment works without any change in the test
|
||||
|
||||
`e2e_gateway.py` holds `Gateway`, a frozen dataclass that wraps a `Transport` and adds the operations tests reuse: `generate_key` / `delete_key` / `key_info`, `model_info`, the LLM calls `chat` / `chat_stream` / `embed` / `ocr`, the spend read-back `spend_logs`, and the poll helpers `poll_logs_for_key` / `poll_logs_for_request_id` that loop to `poll_timeout` instead of sleeping once. Add a new route as a method here so other suites get it for free
|
||||
`proxy_client.py` holds `ProxyClient`, a frozen dataclass that wraps a `Transport` and adds the operations tests reuse: `generate_key` / `delete_key` / `key_info`, `model_info`, the LLM calls `chat` / `chat_stream` / `embed` / `ocr`, the spend read-back `spend_logs`, and the poll helpers `poll_logs_for_key` / `poll_logs_for_request_id` that loop to `poll_timeout` instead of sleeping once. It is exposed as the session-scoped `proxy` fixture (see tests/e2e/conftest.py), which each suite's `client` fixture depends on and injects. Add a new route as a method here so other suites get it for free
|
||||
|
||||
Each suite provides its own `client` fixture (see `llm_translation/passthrough_client.py`), a frozen dataclass that holds the shared `Gateway` and adds suite-specific routes. Cleanup runs through that same `Gateway`, so whatever keys or customers your test creates get torn down by the `resources` fixture
|
||||
Each suite provides its own `client` fixture (see `llm_translation/passthrough_client.py`), a frozen dataclass that holds the shared `ProxyClient` (as `.proxy`) and adds suite-specific routes. Cleanup runs through that same `ProxyClient`, so whatever keys or customers your test creates get torn down by the `resources` fixture
|
||||
|
||||
Request and response bodies are typed pydantic models in `models.py`; only the fields a test reads are modelled, and nothing passes raw dicts. Outcomes come back as a `Result[R]` tagged union (`Success`, `NetworkError`, `UnauthorizedError`, `RateLimitedError`, `ValidationError`, `UnknownApiError`). Handle them with `match`, or call `unwrap(...)` when a non-success should fail the test. The harness hard-fails and never skips: a test marked `e2e` fails when no proxy answers its liveness probe, and once a request reaches the proxy any wrong behavior is likewise a hard failure, so a missing proxy turns the run red instead of being mistaken for a pass
|
||||
|
||||
|
|
|
|||
|
|
@ -113,7 +113,7 @@ class TestPromptCompression:
|
|||
|
||||
def test_prompt_compression_accumulate_spend(self, key_id, user_id):
|
||||
for _ in range(10):
|
||||
response = self.resources.gateway.post("gemini-2.5-flash", key_id, user_id)
|
||||
response = self.resources.proxy.post("gemini-2.5-flash", key_id, user_id)
|
||||
compressed_value = ...
|
||||
assert response.cost == compressed_value # the cost was actually reduced
|
||||
```
|
||||
|
|
@ -128,9 +128,9 @@ The shape is layered so tests stay declarative
|
|||
|
||||
`transport.py` exposes a `Transport` Protocol with `post`, `get`, `delete`, `send`, `stream`, `probe`, plus `bearer(key)` and the `master` header. `HttpTransport` fulfils it, and `SplitTransport` routes each call by path to the data plane or the control plane so a split control-plane/data-plane deployment works without any change in the test
|
||||
|
||||
`e2e_gateway.py` holds `Gateway`, a frozen dataclass that wraps a `Transport` and adds the operations tests reuse: `generate_key` / `delete_key` / `key_info`, `model_info`, the LLM calls `chat` / `chat_stream` / `embed` / `ocr`, the spend read-back `spend_logs`, and the poll helpers `poll_logs_for_key` / `poll_logs_for_request_id` that loop to `poll_timeout` instead of sleeping once. Add a new route as a method here so other suites get it for free
|
||||
`proxy_client.py` holds `ProxyClient`, a frozen dataclass that wraps a `Transport` and adds the operations tests reuse: `generate_key` / `delete_key` / `key_info`, `model_info`, the LLM calls `chat` / `chat_stream` / `embed` / `ocr`, the spend read-back `spend_logs`, and the poll helpers `poll_logs_for_key` / `poll_logs_for_request_id` that loop to `poll_timeout` instead of sleeping once. It is exposed as the session-scoped `proxy` fixture (see tests/e2e/conftest.py), which each suite's `client` fixture depends on and injects. Add a new route as a method here so other suites get it for free
|
||||
|
||||
Each suite provides its own `client` fixture (see `llm_translation/passthrough_client.py`), a frozen dataclass that holds the shared `Gateway` and adds suite-specific routes. Cleanup runs through that same `Gateway`, so whatever keys or customers your test creates get torn down by the `resources` fixture
|
||||
Each suite provides its own `client` fixture (see `llm_translation/passthrough_client.py`), a frozen dataclass that holds the shared `ProxyClient` (as `.proxy`) and adds suite-specific routes. Cleanup runs through that same `ProxyClient`, so whatever keys or customers your test creates get torn down by the `resources` fixture
|
||||
|
||||
Request and response bodies are typed pydantic models in `models.py`; only the fields a test reads are modelled, and nothing passes raw dicts. Outcomes come back as a `Result[R]` tagged union (`Success`, `NetworkError`, `UnauthorizedError`, `RateLimitedError`, `ValidationError`, `UnknownApiError`). Handle them with `match`, or call `unwrap(...)` when a non-success should fail the test. The harness hard-fails and never skips: a test marked `e2e` fails when no proxy answers its liveness probe, and once a request reaches the proxy any wrong behavior is likewise a hard failure, so a missing proxy turns the run red instead of being mistaken for a pass
|
||||
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ from __future__ import annotations
|
|||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from e2e_gateway import Gateway, build_gateway
|
||||
from proxy_client import ProxyClient
|
||||
from e2e_http import StreamingResponse
|
||||
from models import (
|
||||
ChatBody,
|
||||
|
|
@ -21,29 +21,29 @@ ROUTE_NOT_ALLOWED_MARKER = "not allowed to call this route"
|
|||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AccessControlClient:
|
||||
gateway: Gateway
|
||||
proxy: ProxyClient
|
||||
|
||||
def llm_only_key(self) -> str:
|
||||
return self.gateway.generate_key(
|
||||
return self.proxy.generate_key(
|
||||
KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"])
|
||||
)
|
||||
|
||||
def delete_key(self, key: str) -> None:
|
||||
self.gateway.delete_key(key)
|
||||
self.proxy.delete_key(key)
|
||||
|
||||
def chat_status(self, key: str, model: str, content: str) -> StreamingResponse:
|
||||
return self.gateway.transport.send(
|
||||
return self.proxy.transport.send(
|
||||
"/chat/completions",
|
||||
headers=self.gateway.transport.bearer(key),
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=ChatBody(
|
||||
model=model, messages=[ChatMessage(role="user", content=content)]
|
||||
),
|
||||
)
|
||||
|
||||
def create_model_status(self, key: str, model_name: str) -> StreamingResponse:
|
||||
return self.gateway.transport.send(
|
||||
return self.proxy.transport.send(
|
||||
"/model/new",
|
||||
headers=self.gateway.transport.bearer(key),
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=ModelNewBody(
|
||||
model_name=model_name,
|
||||
litellm_params=LiteLLMParamsBody(model="openai/gpt-4o-mini"),
|
||||
|
|
@ -52,5 +52,5 @@ class AccessControlClient:
|
|||
)
|
||||
|
||||
|
||||
def build_client() -> AccessControlClient:
|
||||
return AccessControlClient(gateway=build_gateway())
|
||||
def build_client(proxy: ProxyClient) -> AccessControlClient:
|
||||
return AccessControlClient(proxy=proxy)
|
||||
|
|
|
|||
|
|
@ -3,8 +3,9 @@
|
|||
import pytest
|
||||
|
||||
from access_control_client import AccessControlClient, build_client
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def client() -> AccessControlClient:
|
||||
return build_client()
|
||||
def client(proxy: ProxyClient) -> AccessControlClient:
|
||||
return build_client(proxy)
|
||||
|
|
|
|||
|
|
@ -68,7 +68,7 @@ File delete asserts `object=="file"` and `deleted==True`.
|
|||
|
||||
| File | Covers |
|
||||
|------|--------|
|
||||
| `batch_client.py` | typed file upload/download + batch create/retrieve/cancel/list/delete over the shared Gateway; runtime batch model registration via /model/new; denial helpers |
|
||||
| `batch_client.py` | typed file upload/download + batch create/retrieve/cancel/list/delete over the shared ProxyClient; runtime batch model registration via /model/new; denial helpers |
|
||||
| `capabilities.py` | the provider x scenario matrix + per-provider /model/new params + id-shape classifiers + per-provider raw-id assertion |
|
||||
| `conftest.py` | session-scoped batch deployment registration and teardown |
|
||||
| `test_batches_e2e.py` | parametrized lifecycle with per-endpoint output assertions, file upload/delete outputs, key-model-access denial |
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
"""Client for the batches e2e suite: file upload/download and the batch
|
||||
operations (create / retrieve / cancel / list) over the shared Gateway.
|
||||
operations (create / retrieve / cancel / list) over the shared ProxyClient.
|
||||
|
||||
Batch deployments are registered at runtime via /model/new (see conftest.py),
|
||||
not baked into the proxy config. `create_batch` returns the raw HTTP outcome
|
||||
|
|
@ -16,7 +16,7 @@ from dataclasses import dataclass
|
|||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from e2e_gateway import Gateway, build_gateway
|
||||
from proxy_client import ProxyClient
|
||||
from e2e_http import (
|
||||
FileUploadForm,
|
||||
NoBody,
|
||||
|
|
@ -85,13 +85,13 @@ def is_result_access_denied[R: BaseModel](result: Result[R]) -> bool:
|
|||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BatchClient:
|
||||
gateway: Gateway
|
||||
proxy: ProxyClient
|
||||
|
||||
def create_model(self, model_name: str, litellm_params: LiteLLMParamsBody) -> str:
|
||||
return self.gateway.create_model(model_name, litellm_params, mode="batch")
|
||||
return self.proxy.create_model(model_name, litellm_params, mode="batch")
|
||||
|
||||
def delete_model(self, model_id: str) -> None:
|
||||
self.gateway.delete_model(model_id)
|
||||
self.proxy.delete_model(model_id)
|
||||
|
||||
def upload_file(
|
||||
self,
|
||||
|
|
@ -102,9 +102,9 @@ class BatchClient:
|
|||
model: str | None = None,
|
||||
provider: str | None = None,
|
||||
) -> Result[FileObject]:
|
||||
return self.gateway.transport.upload(
|
||||
return self.proxy.transport.upload(
|
||||
_files_path(provider),
|
||||
headers=self.gateway.transport.bearer(key),
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
form=form,
|
||||
filename="batch_input.jsonl",
|
||||
content=content,
|
||||
|
|
@ -115,18 +115,18 @@ class BatchClient:
|
|||
def create_batch(
|
||||
self, *, body: BatchCreateBody, key: str, provider: str | None = None
|
||||
) -> StreamingResponse:
|
||||
return self.gateway.transport.send(
|
||||
return self.proxy.transport.send(
|
||||
_batches_path(provider),
|
||||
headers=self.gateway.transport.bearer(key),
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=body,
|
||||
)
|
||||
|
||||
def retrieve_batch(
|
||||
self, batch_id: str, *, key: str, provider: str | None = None
|
||||
) -> Result[BatchObject]:
|
||||
return self.gateway.transport.get(
|
||||
return self.proxy.transport.get(
|
||||
f"{_batches_path(provider)}/{batch_id}",
|
||||
headers=self.gateway.transport.bearer(key),
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
params=NoBody(),
|
||||
response_type=BatchObject,
|
||||
)
|
||||
|
|
@ -134,9 +134,9 @@ class BatchClient:
|
|||
def cancel_batch(
|
||||
self, batch_id: str, *, key: str, provider: str | None = None
|
||||
) -> Result[BatchObject]:
|
||||
return self.gateway.transport.post(
|
||||
return self.proxy.transport.post(
|
||||
f"{_batches_path(provider)}/{batch_id}/cancel",
|
||||
headers=self.gateway.transport.bearer(key),
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=NoBody(),
|
||||
response_type=BatchObject,
|
||||
)
|
||||
|
|
@ -144,9 +144,9 @@ class BatchClient:
|
|||
def list_batches(
|
||||
self, *, key: str, provider: str | None = None
|
||||
) -> Result[BatchList]:
|
||||
return self.gateway.transport.get(
|
||||
return self.proxy.transport.get(
|
||||
_batches_path(provider),
|
||||
headers=self.gateway.transport.bearer(key),
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
params=NoBody(),
|
||||
response_type=BatchList,
|
||||
)
|
||||
|
|
@ -154,9 +154,9 @@ class BatchClient:
|
|||
def delete_file(
|
||||
self, file_id: str, *, key: str, provider: str | None = None
|
||||
) -> Result[FileDeleteResponse]:
|
||||
return self.gateway.transport.delete(
|
||||
return self.proxy.transport.delete(
|
||||
f"{_files_path(provider)}/{file_id}",
|
||||
headers=self.gateway.transport.bearer(key),
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=NoBody(),
|
||||
response_type=FileDeleteResponse,
|
||||
)
|
||||
|
|
@ -170,5 +170,5 @@ def _batches_path(provider: str | None) -> str:
|
|||
return f"/{provider}/v1/batches" if provider else "/v1/batches"
|
||||
|
||||
|
||||
def build_client() -> BatchClient:
|
||||
return BatchClient(gateway=build_gateway())
|
||||
def build_client(proxy: ProxyClient) -> BatchClient:
|
||||
return BatchClient(proxy=proxy)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
"""Batches suite's `client` fixture.
|
||||
|
||||
The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker
|
||||
live in the parent tests/e2e/conftest.py. BatchClient holds the shared Gateway, so
|
||||
live in the parent tests/e2e/conftest.py. BatchClient holds the shared ProxyClient, so
|
||||
the `resources` fixture cleans up keys through it; tests register file deletes and
|
||||
batch cancels via `resources.defer(...)`.
|
||||
|
||||
|
|
@ -19,6 +19,7 @@ import pytest
|
|||
from batch_client import BatchClient, build_client
|
||||
from capabilities import PROVIDERS
|
||||
from e2e_http import NoBody
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
|
||||
def pytest_configure(config: pytest.Config) -> None:
|
||||
|
|
@ -29,13 +30,13 @@ def pytest_configure(config: pytest.Config) -> None:
|
|||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def client() -> BatchClient:
|
||||
return build_client()
|
||||
def client(proxy: ProxyClient) -> BatchClient:
|
||||
return build_client(proxy)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def batch_deployments(client: BatchClient) -> Iterator[None]:
|
||||
probe = client.gateway.probe("/health/liveliness", params=NoBody())
|
||||
probe = client.proxy.probe("/health/liveliness", params=NoBody())
|
||||
if not probe.healthy:
|
||||
yield
|
||||
return
|
||||
|
|
|
|||
|
|
@ -372,17 +372,17 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row(
|
|||
environment and OOMed the e2e runner on stage.
|
||||
"""
|
||||
user_id = f"e2e-batch-rl-{unique_marker()}"
|
||||
key = client.gateway.generate_key(
|
||||
key = client.proxy.generate_key(
|
||||
KeyGenerateBody(models=[], tpm_limit=1_000_000, rpm_limit=1_000, user_id=user_id)
|
||||
)
|
||||
resources.defer(lambda: client.gateway.delete_key(key))
|
||||
resources.defer(lambda: client.proxy.delete_key(key))
|
||||
|
||||
window_start = datetime.now(timezone.utc) - timedelta(hours=1)
|
||||
window_end = window_start + timedelta(hours=2)
|
||||
before = frozenset(
|
||||
row.request_id
|
||||
for row in unattributed_rows(
|
||||
client.gateway.spend_logs_window(start=window_start, end=window_end)
|
||||
client.proxy.spend_logs_window(start=window_start, end=window_end)
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -401,12 +401,12 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row(
|
|||
batch = BatchObject.model_validate_json(created.body)
|
||||
resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key)))
|
||||
|
||||
_ = client.gateway.poll_logs_for_key(key, min_rows=1)
|
||||
_ = client.proxy.poll_logs_for_key(key, min_rows=1)
|
||||
|
||||
new_orphans = [
|
||||
row
|
||||
for row in unattributed_rows(
|
||||
client.gateway.spend_logs_window(start=window_start, end=window_end)
|
||||
client.proxy.spend_logs_window(start=window_start, end=window_end)
|
||||
)
|
||||
if row.request_id not in before
|
||||
]
|
||||
|
|
|
|||
|
|
@ -577,30 +577,30 @@ from claude_code._compat_models import ( # noqa: E402
|
|||
)
|
||||
|
||||
|
||||
def _build_control_gateway(proxy: ProxyConfig):
|
||||
def _build_control_plane_client(proxy_config: ProxyConfig):
|
||||
"""Local import of the shared harness so the pure-unit-test tree
|
||||
under ``_driver_unit_tests/`` etc. never has to pull it in. The
|
||||
control plane transport is what /model/new lives on; SplitTransport
|
||||
routes it correctly for both monolithic and split deployments.
|
||||
|
||||
The endpoints come from the *resolved* proxy, not from a second
|
||||
The endpoints come from the *resolved* proxy config, not from a second
|
||||
independent env read, so registration and the cells always hit the
|
||||
same host and key. Both planes get the one URL the cells use; the
|
||||
deployment is fronted by a single address that routes management
|
||||
and LLM paths itself."""
|
||||
from e2e_gateway import build_gateway
|
||||
from proxy_client import build_proxy_client
|
||||
|
||||
return build_gateway(
|
||||
base_url=proxy.base_url,
|
||||
master_key=proxy.api_key,
|
||||
control_plane_base_url=proxy.base_url,
|
||||
return build_proxy_client(
|
||||
base_url=proxy_config.base_url,
|
||||
master_key=proxy_config.api_key,
|
||||
control_plane_base_url=proxy_config.base_url,
|
||||
)
|
||||
|
||||
|
||||
def _register_deployment(gateway, deployment: CompatDeployment) -> str:
|
||||
def _register_deployment(proxy, deployment: CompatDeployment) -> str:
|
||||
"""Register one deployment and return its proxy-assigned model_id
|
||||
once it is servable on the data plane."""
|
||||
return gateway.create_model(
|
||||
return proxy.create_model(
|
||||
deployment.model_name,
|
||||
deployment.litellm_params,
|
||||
)
|
||||
|
|
@ -624,20 +624,20 @@ def _compat_models_registered() -> Any:
|
|||
but do not abort the session: the cells that need that specific
|
||||
deployment will 400 with "Invalid model name" and fail loudly,
|
||||
which is the right signal (missing cred on the proxy side)."""
|
||||
proxy = resolve_proxy()
|
||||
if proxy is None:
|
||||
proxy_config = resolve_proxy()
|
||||
if proxy_config is None:
|
||||
yield
|
||||
return
|
||||
|
||||
from requests import RequestException
|
||||
|
||||
gateway = _build_control_gateway(proxy)
|
||||
proxy = _build_control_plane_client(proxy_config)
|
||||
registered_ids: list[str] = []
|
||||
failures: list[tuple[str, str]] = []
|
||||
try:
|
||||
for deployment in load_all_deployments():
|
||||
try:
|
||||
model_id = _register_deployment(gateway, deployment)
|
||||
model_id = _register_deployment(proxy, deployment)
|
||||
registered_ids.append(model_id)
|
||||
except (AssertionError, RequestException) as exc:
|
||||
failures.append((deployment.model_name, str(exc)))
|
||||
|
|
@ -656,7 +656,7 @@ def _compat_models_registered() -> Any:
|
|||
finally:
|
||||
for model_id in registered_ids:
|
||||
try:
|
||||
gateway.delete_model(model_id)
|
||||
proxy.delete_model(model_id)
|
||||
except (AssertionError, RequestException):
|
||||
# Best-effort — teardown surfaces via warnings inside
|
||||
# ``delete_model`` already; swallowing here so one flaky
|
||||
|
|
|
|||
|
|
@ -23,7 +23,8 @@ import requests
|
|||
|
||||
from e2e_config import CONTROL_PLANE_BASE_URL, PROXY_BASE_URL
|
||||
from junit_properties import attach_result_properties
|
||||
from lifecycle import GatewayProvider, ResourceManager
|
||||
from lifecycle import ProxyClientProvider, ResourceManager
|
||||
from proxy_client import ProxyClient, build_proxy_client
|
||||
|
||||
|
||||
_E2E_TEST_RAN = pytest.StashKey[bool]()
|
||||
|
|
@ -120,11 +121,18 @@ def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None:
|
|||
sys.path.remove(spend_dir)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def proxy() -> ProxyClient:
|
||||
"""The shared ProxyClient every suite's client is built from. Suite `client`
|
||||
fixtures depend on this and inject it, so the proxy wiring lives in one place."""
|
||||
return build_proxy_client()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def resources(client: GatewayProvider) -> Iterator[ResourceManager]:
|
||||
def resources(client: ProxyClientProvider) -> Iterator[ResourceManager]:
|
||||
"""init -> run -> teardown: create a manager, run the test, release resources.
|
||||
Cleanup goes through the shared Gateway, whatever the suite's client adds."""
|
||||
manager = ResourceManager(client=client.gateway)
|
||||
Cleanup goes through the shared ProxyClient, whatever the suite's client adds."""
|
||||
manager = ResourceManager(client=client.proxy)
|
||||
manager.init()
|
||||
yield manager
|
||||
manager.teardown()
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ 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 e2e_gateway import Gateway
|
||||
from proxy_client import ProxyClient
|
||||
from models import KeyGenerateBody
|
||||
|
||||
|
||||
|
|
@ -52,7 +52,7 @@ def run_case(case: E2ECase) -> None:
|
|||
@runtime_checkable
|
||||
class ResourceClient(Protocol):
|
||||
"""Proxy operations the convenience creators use. Resource types without a
|
||||
creator here are handled generically via ResourceManager.defer(). The Gateway
|
||||
creator here are handled generically via ResourceManager.defer(). The ProxyClient
|
||||
satisfies this."""
|
||||
|
||||
def generate_key(self, body: KeyGenerateBody) -> str: ...
|
||||
|
|
@ -63,12 +63,12 @@ class ResourceClient(Protocol):
|
|||
|
||||
|
||||
@runtime_checkable
|
||||
class GatewayProvider(Protocol):
|
||||
"""Every suite's client exposes the shared Gateway, which the resources fixture
|
||||
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 gateway(self) -> Gateway: ...
|
||||
def proxy(self) -> ProxyClient: ...
|
||||
|
||||
|
||||
@dataclass
|
||||
|
|
|
|||
|
|
@ -2,13 +2,14 @@
|
|||
|
||||
The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker
|
||||
live in the parent tests/e2e/conftest.py. PassthroughClient holds the shared
|
||||
Gateway, so the `resources` fixture cleans up keys this suite creates.
|
||||
ProxyClient, so the `resources` fixture cleans up keys this suite creates.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from endpoints_client import EndpointsClient, build_endpoints_client
|
||||
from passthrough_client import PassthroughClient, build_client
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
|
||||
def pytest_configure(config: pytest.Config) -> None:
|
||||
|
|
@ -19,10 +20,10 @@ def pytest_configure(config: pytest.Config) -> None:
|
|||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def client() -> PassthroughClient:
|
||||
return build_client()
|
||||
def client(proxy: ProxyClient) -> PassthroughClient:
|
||||
return build_client(proxy)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def endpoints_client() -> EndpointsClient:
|
||||
return build_endpoints_client()
|
||||
def endpoints_client(proxy: ProxyClient) -> EndpointsClient:
|
||||
return build_endpoints_client(proxy)
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ from dataclasses import dataclass
|
|||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from e2e_gateway import Gateway, build_gateway
|
||||
from proxy_client import ProxyClient
|
||||
from e2e_http import StreamingResponse
|
||||
from models import ChatMessage, LiteLLMParamsBody
|
||||
|
||||
|
|
@ -156,17 +156,17 @@ class ImagesResult(BaseModel):
|
|||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class EndpointsClient:
|
||||
gateway: Gateway
|
||||
proxy: ProxyClient
|
||||
|
||||
def create_model(self, model_name: str, litellm_params: LiteLLMParamsBody) -> str:
|
||||
return self.gateway.create_model(model_name, litellm_params)
|
||||
return self.proxy.create_model(model_name, litellm_params)
|
||||
|
||||
def delete_model(self, model_id: str) -> None:
|
||||
self.gateway.delete_model(model_id)
|
||||
self.proxy.delete_model(model_id)
|
||||
|
||||
def _send(self, path: str, key: str, body: BaseModel) -> StreamingResponse:
|
||||
return self.gateway.transport.send(
|
||||
path, headers=self.gateway.transport.bearer(key), json=body
|
||||
return self.proxy.transport.send(
|
||||
path, headers=self.proxy.transport.bearer(key), json=body
|
||||
)
|
||||
|
||||
def responses(self, key: str, model: str, text: str) -> StreamingResponse:
|
||||
|
|
@ -216,5 +216,5 @@ class EndpointsClient:
|
|||
)
|
||||
|
||||
|
||||
def build_endpoints_client() -> EndpointsClient:
|
||||
return EndpointsClient(gateway=build_gateway())
|
||||
def build_endpoints_client(proxy: ProxyClient) -> EndpointsClient:
|
||||
return EndpointsClient(proxy=proxy)
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ from dataclasses import dataclass
|
|||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from e2e_gateway import Gateway, build_gateway
|
||||
from proxy_client import ProxyClient
|
||||
from e2e_http import Headers, StreamingResponse
|
||||
from models import ChatMessage
|
||||
|
||||
|
|
@ -108,7 +108,7 @@ def _tags_header(tags: list[str] | None) -> str | None:
|
|||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PassthroughClient:
|
||||
gateway: Gateway
|
||||
proxy: ProxyClient
|
||||
|
||||
# ---- Gemini native passthrough (/gemini/v1beta/...) -----------------
|
||||
|
||||
|
|
@ -121,7 +121,7 @@ class PassthroughClient:
|
|||
tools: list[GeminiTool] | None = None,
|
||||
tags: list[str] | None = None,
|
||||
) -> StreamingResponse:
|
||||
return self.gateway.transport.send(
|
||||
return self.proxy.transport.send(
|
||||
f"/gemini/v1beta/models/{model}:generateContent",
|
||||
headers=GeminiHeaders(x_goog_api_key=key, tags=_tags_header(tags)),
|
||||
json=GeminiGenerateBody(
|
||||
|
|
@ -132,7 +132,7 @@ class PassthroughClient:
|
|||
def gemini_stream(
|
||||
self, key: str, model: str, text: str, *, tags: list[str] | None = None
|
||||
) -> StreamingResponse:
|
||||
return self.gateway.transport.send(
|
||||
return self.proxy.transport.send(
|
||||
f"/gemini/v1beta/models/{model}:streamGenerateContent",
|
||||
headers=GeminiHeaders(x_goog_api_key=key, tags=_tags_header(tags)),
|
||||
json=GeminiGenerateBody(
|
||||
|
|
@ -151,7 +151,7 @@ class PassthroughClient:
|
|||
f"/vertex_ai/v1/projects/{project}/locations/{location}"
|
||||
f"/publishers/google/models/{model}:generateContent"
|
||||
)
|
||||
return self.gateway.transport.send(
|
||||
return self.proxy.transport.send(
|
||||
path,
|
||||
headers=VertexHeaders(x_litellm_api_key=key),
|
||||
json=GeminiGenerateBody(
|
||||
|
|
@ -172,7 +172,7 @@ class PassthroughClient:
|
|||
stream: bool = False,
|
||||
tags: list[str] | None = None,
|
||||
) -> StreamingResponse:
|
||||
return self.gateway.transport.send(
|
||||
return self.proxy.transport.send(
|
||||
"/anthropic/v1/messages",
|
||||
headers=AnthropicHeaders(x_api_key=key, tags=_tags_header(tags)),
|
||||
json=AnthropicMessageBody(
|
||||
|
|
@ -186,5 +186,5 @@ class PassthroughClient:
|
|||
)
|
||||
|
||||
|
||||
def build_client() -> PassthroughClient:
|
||||
return PassthroughClient(gateway=build_gateway())
|
||||
def build_client(proxy: ProxyClient) -> PassthroughClient:
|
||||
return PassthroughClient(proxy=proxy)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
"""Realtime suite's `client` and `realtime_models` fixtures.
|
||||
|
||||
The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker
|
||||
live in the parent tests/e2e/conftest.py. RealtimeClient holds the shared Gateway,
|
||||
live in the parent tests/e2e/conftest.py. RealtimeClient holds the shared ProxyClient,
|
||||
so the `resources` fixture cleans up keys this suite creates.
|
||||
|
||||
`realtime_models` registers every provider's realtime deployment through /model/new
|
||||
|
|
@ -15,11 +15,12 @@ from collections.abc import Iterator
|
|||
import pytest
|
||||
|
||||
from realtime_client import PROVIDERS, RealtimeClient, build_client
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def client() -> RealtimeClient:
|
||||
return build_client()
|
||||
def client(proxy: ProxyClient) -> RealtimeClient:
|
||||
return build_client(proxy)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
|
|
@ -34,4 +35,4 @@ def realtime_models(client: RealtimeClient) -> Iterator[dict[str, str]]:
|
|||
yield {provider_id: model_name for provider_id, model_name, _ in records}
|
||||
finally:
|
||||
for _, _, model_id in records:
|
||||
client.gateway.delete_model(model_id)
|
||||
client.proxy.delete_model(model_id)
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ from websockets.sync.client import connect
|
|||
from websockets.sync.connection import Connection
|
||||
|
||||
from e2e_config import PROXY_BASE_URL, unique_marker
|
||||
from e2e_gateway import Gateway, build_gateway
|
||||
from proxy_client import ProxyClient
|
||||
from models import LiteLLMParamsBody
|
||||
|
||||
_M = TypeVar("_M", bound=BaseModel)
|
||||
|
|
@ -329,7 +329,7 @@ class RealtimeSession:
|
|||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RealtimeClient:
|
||||
gateway: Gateway
|
||||
proxy: ProxyClient
|
||||
|
||||
def provision(self, provider: RealtimeProvider) -> tuple[str, str]:
|
||||
"""Register this provider's realtime deployment through /model/new and return
|
||||
|
|
@ -338,7 +338,7 @@ class RealtimeClient:
|
|||
show up as a realtime model on /model/info. add_deployment runs synchronously,
|
||||
so the deployment is connectable as soon as this returns."""
|
||||
model_name = f"{provider.alias}-{unique_marker()}"
|
||||
model_id = self.gateway.create_model(
|
||||
model_id = self.proxy.create_model(
|
||||
model_name, provider.litellm_params, mode="realtime"
|
||||
)
|
||||
return model_name, model_id
|
||||
|
|
@ -355,5 +355,5 @@ class RealtimeClient:
|
|||
yield RealtimeSession(connection=connection)
|
||||
|
||||
|
||||
def build_client() -> RealtimeClient:
|
||||
return RealtimeClient(gateway=build_gateway())
|
||||
def build_client(proxy: ProxyClient) -> RealtimeClient:
|
||||
return RealtimeClient(proxy=proxy)
|
||||
|
|
|
|||
|
|
@ -81,9 +81,9 @@ def _cache_chat(
|
|||
RichMessage(role="user", content=[TextBlock(text="Reply with one word.")]),
|
||||
],
|
||||
)
|
||||
return client.gateway.transport.post(
|
||||
return client.proxy.transport.post(
|
||||
"/chat/completions",
|
||||
headers=client.gateway.transport.bearer(key),
|
||||
headers=client.proxy.transport.bearer(key),
|
||||
json=body,
|
||||
response_type=ChatResponse,
|
||||
)
|
||||
|
|
@ -120,11 +120,11 @@ class TestCacheControl:
|
|||
self, client: PassthroughClient, resources: ResourceManager
|
||||
) -> None:
|
||||
model = f"e2e-bedrock-cache-{unique_marker()}"
|
||||
model_id = client.gateway.create_model(
|
||||
model_id = client.proxy.create_model(
|
||||
model,
|
||||
LiteLLMParamsBody(model=BEDROCK_MODEL, aws_region_name="us-east-1"),
|
||||
)
|
||||
resources.defer(lambda: client.gateway.delete_model(model_id))
|
||||
resources.defer(lambda: client.proxy.delete_model(model_id))
|
||||
_assert_cache_read_on_second_call(client, resources.key(), model)
|
||||
|
||||
@pytest.mark.covers(
|
||||
|
|
@ -135,7 +135,7 @@ class TestCacheControl:
|
|||
self, client: PassthroughClient, resources: ResourceManager
|
||||
) -> None:
|
||||
model = f"e2e-vertex-cache-{unique_marker()}"
|
||||
model_id = client.gateway.create_model(
|
||||
model_id = client.proxy.create_model(
|
||||
model,
|
||||
LiteLLMParamsBody(
|
||||
model=VERTEX_MODEL,
|
||||
|
|
@ -144,5 +144,5 @@ class TestCacheControl:
|
|||
vertex_credentials=os.environ.get("VERTEXAI_CREDENTIALS"),
|
||||
),
|
||||
)
|
||||
resources.defer(lambda: client.gateway.delete_model(model_id))
|
||||
resources.defer(lambda: client.proxy.delete_model(model_id))
|
||||
_assert_cache_read_on_second_call(client, resources.key(), model)
|
||||
|
|
|
|||
|
|
@ -43,7 +43,7 @@ class TestChatCompletionsRegression:
|
|||
self, client: PassthroughClient, scoped_key: str, model: str, route: str
|
||||
) -> None:
|
||||
response = unwrap(
|
||||
client.gateway.chat(
|
||||
client.proxy.chat(
|
||||
scoped_key,
|
||||
ChatBody(
|
||||
model=model,
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ import pytest
|
|||
from pydantic import BaseModel, RootModel
|
||||
|
||||
from e2e_config import unique_marker
|
||||
from e2e_gateway import Gateway
|
||||
from proxy_client import ProxyClient
|
||||
from e2e_http import Success, unwrap
|
||||
from endpoints_client import EndpointsClient
|
||||
from lifecycle import ResourceManager
|
||||
|
|
@ -116,14 +116,14 @@ def _model_info_entry(entries: list[ModelInfoEntry], model_name: str) -> ModelIn
|
|||
pytest.fail(f"{model_name} absent from /model/info; the override did not load")
|
||||
|
||||
|
||||
def _poll_breakdown_row(gateway: Gateway, key: str, response_id: str | None) -> _SpendRow:
|
||||
def _poll_breakdown_row(proxy: ProxyClient, key: str, response_id: str | None) -> _SpendRow:
|
||||
"""Poll /spend/logs until the call's row lands with a cost breakdown (rows
|
||||
flush ~60s behind the call via proxy_batch_write_at)."""
|
||||
deadline = time.monotonic() + gateway.poll_timeout
|
||||
deadline = time.monotonic() + proxy.poll_timeout
|
||||
while time.monotonic() < deadline:
|
||||
result = gateway.transport.get(
|
||||
result = proxy.transport.get(
|
||||
"/spend/logs",
|
||||
headers=gateway.transport.master,
|
||||
headers=proxy.transport.master,
|
||||
params=SpendLogsParams(api_key=key),
|
||||
response_type=_SpendRows,
|
||||
)
|
||||
|
|
@ -144,7 +144,7 @@ def _poll_breakdown_row(gateway: Gateway, key: str, response_id: str | None) ->
|
|||
return row
|
||||
if priced and response_id is None:
|
||||
return priced[0]
|
||||
time.sleep(gateway.poll_interval)
|
||||
time.sleep(proxy.poll_interval)
|
||||
pytest.fail("no spend row with a cost breakdown landed before the deadline")
|
||||
|
||||
|
||||
|
|
@ -158,7 +158,7 @@ class TestCustomPricing:
|
|||
model = _provision_custom_priced(endpoints_client, resources)
|
||||
|
||||
chat = unwrap(
|
||||
endpoints_client.gateway.chat(
|
||||
endpoints_client.proxy.chat(
|
||||
scoped_key,
|
||||
ChatBody(
|
||||
model=model,
|
||||
|
|
@ -172,7 +172,7 @@ class TestCustomPricing:
|
|||
)
|
||||
)
|
||||
|
||||
row = _poll_breakdown_row(endpoints_client.gateway, scoped_key, chat.id)
|
||||
row = _poll_breakdown_row(endpoints_client.proxy, scoped_key, chat.id)
|
||||
assert row.metadata and row.metadata.cost_breakdown # guaranteed by the poll
|
||||
breakdown = row.metadata.cost_breakdown
|
||||
|
||||
|
|
@ -198,7 +198,7 @@ class TestCustomPricing:
|
|||
self, endpoints_client: EndpointsClient, resources: ResourceManager
|
||||
) -> None:
|
||||
model = _provision_custom_priced(endpoints_client, resources)
|
||||
entry = _model_info_entry(endpoints_client.gateway.model_info(), model)
|
||||
entry = _model_info_entry(endpoints_client.proxy.model_info(), model)
|
||||
|
||||
assert entry.litellm_params.input_cost_per_token == CUSTOM_INPUT_RATE, (
|
||||
f"/model/info litellm_params input rate "
|
||||
|
|
@ -223,7 +223,7 @@ class TestCustomPricing:
|
|||
output_cost_per_token=None,
|
||||
)
|
||||
|
||||
entries = {entry.model_name: entry for entry in endpoints_client.gateway.model_info()}
|
||||
entries = {entry.model_name: entry for entry in endpoints_client.proxy.model_info()}
|
||||
custom_entry = entries.get(custom)
|
||||
sibling_entry = entries.get(sibling)
|
||||
assert custom_entry is not None, f"{custom} absent from /model/info"
|
||||
|
|
|
|||
|
|
@ -33,11 +33,11 @@ PROMPT = "What is 17 + 26? Answer with just the number."
|
|||
|
||||
def _register_reasoner(client: PassthroughClient, resources: ResourceManager) -> str:
|
||||
model = f"e2e-deepseek-reasoner-{unique_marker()}"
|
||||
model_id = client.gateway.create_model(
|
||||
model_id = client.proxy.create_model(
|
||||
model,
|
||||
LiteLLMParamsBody(model=REASONER, api_key="os.environ/DEEPSEEK_API_KEY"),
|
||||
)
|
||||
resources.defer(lambda: client.gateway.delete_model(model_id))
|
||||
resources.defer(lambda: client.proxy.delete_model(model_id))
|
||||
return model
|
||||
|
||||
|
||||
|
|
@ -56,7 +56,7 @@ class TestDeepSeekReasoningDisable:
|
|||
key = resources.key()
|
||||
|
||||
response = unwrap(
|
||||
client.gateway.chat(
|
||||
client.proxy.chat(
|
||||
key,
|
||||
ChatBody(
|
||||
model=model,
|
||||
|
|
@ -78,7 +78,7 @@ class TestDeepSeekReasoningDisable:
|
|||
key = resources.key()
|
||||
|
||||
response = unwrap(
|
||||
client.gateway.chat(
|
||||
client.proxy.chat(
|
||||
key,
|
||||
ChatBody(
|
||||
model=model,
|
||||
|
|
@ -100,7 +100,7 @@ class TestDeepSeekReasoningDisable:
|
|||
key = resources.key()
|
||||
|
||||
response = unwrap(
|
||||
client.gateway.chat(
|
||||
client.proxy.chat(
|
||||
key,
|
||||
ChatBody(
|
||||
model=model,
|
||||
|
|
|
|||
|
|
@ -76,9 +76,9 @@ def _system_reminder_turn() -> RichMessage:
|
|||
def _post_messages(
|
||||
client: EndpointsClient, key: str, body: RichMessagesRequest
|
||||
) -> Result[MessagesResult]:
|
||||
return client.gateway.transport.post(
|
||||
return client.proxy.transport.post(
|
||||
"/v1/messages",
|
||||
headers=client.gateway.transport.bearer(key),
|
||||
headers=client.proxy.transport.bearer(key),
|
||||
json=body,
|
||||
response_type=MessagesResult,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -150,7 +150,7 @@ class TestRustOcrGateway:
|
|||
resources.defer(lambda: endpoints_client.delete_model(model_id))
|
||||
key = resources.key()
|
||||
|
||||
response = unwrap(endpoints_client.gateway.ocr(key, OcrBody(model=model, document=case.document)))
|
||||
response = unwrap(endpoints_client.proxy.ocr(key, OcrBody(model=model, document=case.document)))
|
||||
_assert_ocr_document(response)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -35,7 +35,7 @@ def _fetch_cost_breakdown(client: PassthroughClient, result: StreamingResponse)
|
|||
whole point of passthrough spend tracking.
|
||||
"""
|
||||
assert result.call_id, "passthrough response had no x-litellm-call-id header"
|
||||
rows = client.gateway.poll_logs_for_request_id(
|
||||
rows = client.proxy.poll_logs_for_request_id(
|
||||
result.call_id,
|
||||
predicate=lambda rs: (rs[0].spend or 0) > 0,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -37,17 +37,17 @@ class TestServiceTier:
|
|||
self, client: PassthroughClient, resources: ResourceManager
|
||||
) -> None:
|
||||
model = f"e2e-service-tier-{unique_marker()}"
|
||||
model_id = client.gateway.create_model(
|
||||
model_id = client.proxy.create_model(
|
||||
model,
|
||||
LiteLLMParamsBody(
|
||||
model="openai/gpt-5.5", api_key="os.environ/OPENAI_API_KEY"
|
||||
),
|
||||
)
|
||||
resources.defer(lambda: client.gateway.delete_model(model_id))
|
||||
resources.defer(lambda: client.proxy.delete_model(model_id))
|
||||
key = resources.key()
|
||||
|
||||
response = unwrap(
|
||||
client.gateway.chat(
|
||||
client.proxy.chat(
|
||||
key,
|
||||
ChatBody(
|
||||
model=model,
|
||||
|
|
|
|||
|
|
@ -91,9 +91,9 @@ def _add_vertex_passthrough_model(
|
|||
client: PassthroughClient, model_name: str, project: str, credentials: str
|
||||
) -> str:
|
||||
return unwrap(
|
||||
client.gateway.transport.post(
|
||||
client.proxy.transport.post(
|
||||
"/model/new",
|
||||
headers=client.gateway.transport.master,
|
||||
headers=client.proxy.transport.master,
|
||||
json=_ModelNewBody(
|
||||
model_name=model_name,
|
||||
litellm_params=_VertexDeploymentParams(
|
||||
|
|
@ -111,9 +111,9 @@ def _add_vertex_passthrough_model(
|
|||
|
||||
|
||||
def _delete_model(client: PassthroughClient, model_id: str) -> None:
|
||||
_ = client.gateway.transport.post(
|
||||
_ = client.proxy.transport.post(
|
||||
"/model/delete",
|
||||
headers=client.gateway.transport.master,
|
||||
headers=client.proxy.transport.master,
|
||||
json=_ModelDeleteBody(id=model_id),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -126,7 +126,7 @@ def _costed_row(client: PassthroughClient, call_id: str | None) -> SpendLogRow:
|
|||
a billed Vertex call that LiteLLM did not track is the exact regression #31689
|
||||
guards against."""
|
||||
assert call_id, "vertex passthrough response had no x-litellm-call-id header"
|
||||
rows = client.gateway.poll_logs_for_request_id(
|
||||
rows = client.proxy.poll_logs_for_request_id(
|
||||
call_id,
|
||||
predicate=lambda rs: (rs[0].spend or 0) > 0,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ import pytest
|
|||
from logging_client import LangfuseCreds, LoggingClient, build_logging_client, load_langfuse_creds
|
||||
from datadog_reader import DdLogsReader, build_dd_logs_reader
|
||||
from otel_client import OtelReader, build_otel_reader
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
|
||||
def pytest_configure(config: pytest.Config) -> None:
|
||||
|
|
@ -23,11 +24,11 @@ def pytest_configure(config: pytest.Config) -> None:
|
|||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def client() -> LoggingClient:
|
||||
"""The logging suite's client: holds the shared Gateway so `resources` /
|
||||
def client(proxy: ProxyClient) -> LoggingClient:
|
||||
"""The logging suite's client: holds the shared ProxyClient so `resources` /
|
||||
`scoped_key` clean up keys and teams, and adds `/metrics` scraping plus
|
||||
Langfuse read-back."""
|
||||
return build_logging_client()
|
||||
return build_logging_client(proxy)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
"""Client for the logging e2e suite: team/key/org-scoped Langfuse OTEL callbacks,
|
||||
chat (including tools), Prometheus scrape, and Langfuse observation read-back.
|
||||
|
||||
Holds the shared Gateway so the ``resources`` fixture cleans up keys, teams,
|
||||
Holds the shared ProxyClient so the ``resources`` fixture cleans up keys, teams,
|
||||
users, orgs, and models it creates. External Langfuse reads go through
|
||||
``e2e_http`` (the only module allowed to call ``requests.*``).
|
||||
|
||||
|
|
@ -24,7 +24,7 @@ import pytest
|
|||
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT
|
||||
from e2e_gateway import Gateway, build_gateway
|
||||
from proxy_client import ProxyClient
|
||||
from e2e_http import (
|
||||
URL,
|
||||
AuthHeaders,
|
||||
|
|
@ -262,7 +262,7 @@ def observation_has_guardrail(obs: LangfuseObservation, *, guardrail_name: str)
|
|||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class LoggingClient:
|
||||
gateway: Gateway
|
||||
proxy: ProxyClient
|
||||
|
||||
def key_with_alias(
|
||||
self,
|
||||
|
|
@ -274,7 +274,7 @@ class LoggingClient:
|
|||
organization_id: str | None = None,
|
||||
metadata: KeyMetadata | None = None,
|
||||
) -> str:
|
||||
return self.gateway.generate_key(
|
||||
return self.proxy.generate_key(
|
||||
KeyGenerateBody(
|
||||
key_alias=alias,
|
||||
models=models,
|
||||
|
|
@ -286,7 +286,7 @@ class LoggingClient:
|
|||
)
|
||||
|
||||
def delete_key(self, key: str) -> None:
|
||||
self.gateway.delete_key(key)
|
||||
self.proxy.delete_key(key)
|
||||
|
||||
def create_team(
|
||||
self,
|
||||
|
|
@ -296,9 +296,9 @@ class LoggingClient:
|
|||
organization_id: str | None = None,
|
||||
) -> str:
|
||||
return unwrap(
|
||||
self.gateway.transport.post(
|
||||
self.proxy.transport.post(
|
||||
"/team/new",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=TeamNewBody(
|
||||
team_alias=alias,
|
||||
models=models,
|
||||
|
|
@ -309,18 +309,18 @@ class LoggingClient:
|
|||
).team_id
|
||||
|
||||
def delete_team(self, team_id: str) -> None:
|
||||
_ = self.gateway.transport.post(
|
||||
_ = self.proxy.transport.post(
|
||||
"/team/delete",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=TeamDeleteBody(team_ids=[team_id]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
||||
def create_user(self, *, user_email: str, user_id: str | None = None) -> str:
|
||||
return unwrap(
|
||||
self.gateway.transport.post(
|
||||
self.proxy.transport.post(
|
||||
"/user/new",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=UserNewBody(
|
||||
user_email=user_email,
|
||||
user_role="internal_user",
|
||||
|
|
@ -331,27 +331,27 @@ class LoggingClient:
|
|||
).user_id
|
||||
|
||||
def delete_user(self, user_id: str) -> None:
|
||||
_ = self.gateway.transport.post(
|
||||
_ = self.proxy.transport.post(
|
||||
"/user/delete",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=UserDeleteBody(user_ids=[user_id]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
||||
def create_org(self, alias: str, *, models: list[str]) -> str:
|
||||
return unwrap(
|
||||
self.gateway.transport.post(
|
||||
self.proxy.transport.post(
|
||||
"/organization/new",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=OrgNewBody(organization_alias=alias, models=models),
|
||||
response_type=OrgNewResponse,
|
||||
)
|
||||
).organization_id
|
||||
|
||||
def delete_org(self, organization_id: str) -> None:
|
||||
_ = self.gateway.transport.delete(
|
||||
_ = self.proxy.transport.delete(
|
||||
"/organization/delete",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=OrgDeleteBody(organization_ids=[organization_id]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -364,9 +364,9 @@ class LoggingClient:
|
|||
callback_type: Literal["success", "failure", "success_and_failure"] = "success_and_failure",
|
||||
) -> None:
|
||||
response = unwrap(
|
||||
self.gateway.transport.post(
|
||||
self.proxy.transport.post(
|
||||
f"/team/{team_id}/callback",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=TeamCallbackBody(
|
||||
callback_name="langfuse_otel",
|
||||
callback_type=callback_type,
|
||||
|
|
@ -382,9 +382,9 @@ class LoggingClient:
|
|||
def create_tool_permission_guardrail(self, name: str, *, allowed_tool: str) -> str:
|
||||
"""Register a tool_permission guardrail that allows one tool and denies the rest."""
|
||||
response = unwrap(
|
||||
self.gateway.transport.post(
|
||||
self.proxy.transport.post(
|
||||
"/guardrails",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=CreateGuardrailBody(
|
||||
guardrail=GuardrailSpec(
|
||||
guardrail_name=name,
|
||||
|
|
@ -412,22 +412,22 @@ class LoggingClient:
|
|||
return guardrail_id
|
||||
|
||||
def delete_guardrail(self, guardrail_id: str) -> None:
|
||||
_ = self.gateway.transport.delete(
|
||||
_ = self.proxy.transport.delete(
|
||||
f"/guardrails/{guardrail_id}",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=NoBody(),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
||||
def create_model(self, model_name: str, litellm_params: LiteLLMParamsBody) -> str:
|
||||
return self.gateway.create_model(model_name, litellm_params)
|
||||
return self.proxy.create_model(model_name, litellm_params)
|
||||
|
||||
def delete_model(self, model_id: str) -> None:
|
||||
self.gateway.delete_model(model_id)
|
||||
self.proxy.delete_model(model_id)
|
||||
|
||||
def chat(self, key: str, model: str, text: str) -> ChatResponse:
|
||||
return unwrap(
|
||||
self.gateway.chat(
|
||||
self.proxy.chat(
|
||||
key,
|
||||
ChatBody(
|
||||
model=model,
|
||||
|
|
@ -459,10 +459,10 @@ class LoggingClient:
|
|||
guardrails=guardrails,
|
||||
)
|
||||
if stream:
|
||||
return self.gateway.chat_stream(key, body)
|
||||
return self.gateway.transport.send(
|
||||
return self.proxy.chat_stream(key, body)
|
||||
return self.proxy.transport.send(
|
||||
"/chat/completions",
|
||||
headers=self.gateway.transport.bearer(key),
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=body,
|
||||
)
|
||||
|
||||
|
|
@ -479,11 +479,11 @@ class LoggingClient:
|
|||
stream=True if stream else None,
|
||||
)
|
||||
if stream:
|
||||
return self.gateway.transport.stream(
|
||||
"/v1/messages", headers=self.gateway.transport.bearer(key), json=body
|
||||
return self.proxy.transport.stream(
|
||||
"/v1/messages", headers=self.proxy.transport.bearer(key), json=body
|
||||
)
|
||||
return self.gateway.transport.send(
|
||||
"/v1/messages", headers=self.gateway.transport.bearer(key), json=body
|
||||
return self.proxy.transport.send(
|
||||
"/v1/messages", headers=self.proxy.transport.bearer(key), json=body
|
||||
)
|
||||
|
||||
def responses_raw(
|
||||
|
|
@ -498,15 +498,15 @@ class LoggingClient:
|
|||
model=model, input=text, max_output_tokens=max_output_tokens, stream=True if stream else None
|
||||
)
|
||||
if stream:
|
||||
return self.gateway.transport.stream(
|
||||
"/v1/responses", headers=self.gateway.transport.bearer(key), json=body
|
||||
return self.proxy.transport.stream(
|
||||
"/v1/responses", headers=self.proxy.transport.bearer(key), json=body
|
||||
)
|
||||
return self.gateway.transport.send(
|
||||
"/v1/responses", headers=self.gateway.transport.bearer(key), json=body
|
||||
return self.proxy.transport.send(
|
||||
"/v1/responses", headers=self.proxy.transport.bearer(key), json=body
|
||||
)
|
||||
|
||||
def scrape_metrics(self) -> str:
|
||||
return self.gateway.probe("/metrics", params=NoBody()).body
|
||||
return self.proxy.probe("/metrics", params=NoBody()).body
|
||||
|
||||
def poll_proxy_spend_for_key(
|
||||
self,
|
||||
|
|
@ -529,7 +529,7 @@ class LoggingClient:
|
|||
return False
|
||||
return True
|
||||
|
||||
rows = self.gateway.poll_logs_for_key(
|
||||
rows = self.proxy.poll_logs_for_key(
|
||||
key, min_rows=1, predicate=lambda rs: any(_matches(r) for r in rs)
|
||||
)
|
||||
for row in rows:
|
||||
|
|
@ -623,15 +623,15 @@ def first_ok(client: LoggingClient, send: Callable[[], StreamingResponse]) -> St
|
|||
the data plane's auth cache picks it up, so retry on 401 to a deadline; a
|
||||
401 is rejected before the LLM call, so it cannot contaminate delivery or
|
||||
trace assertions. Any other failure is behavior under test and fails hard."""
|
||||
deadline = time.monotonic() + client.gateway.poll_timeout
|
||||
deadline = time.monotonic() + client.proxy.poll_timeout
|
||||
while True:
|
||||
outcome = send()
|
||||
if outcome.ok:
|
||||
return outcome
|
||||
if outcome.status_code != 401 or time.monotonic() >= deadline:
|
||||
require_successful_call(outcome)
|
||||
time.sleep(client.gateway.poll_interval)
|
||||
time.sleep(client.proxy.poll_interval)
|
||||
|
||||
|
||||
def build_logging_client() -> LoggingClient:
|
||||
return LoggingClient(gateway=build_gateway())
|
||||
def build_logging_client(proxy: ProxyClient) -> LoggingClient:
|
||||
return LoggingClient(proxy=proxy)
|
||||
|
|
|
|||
|
|
@ -52,7 +52,7 @@ def _assert_datadog_configured(client: LoggingClient) -> None:
|
|||
"""Recorded state: the proxy reports the DataDog callback among its active
|
||||
callbacks, so a missing destination config fails here, before any
|
||||
delivery-based assertion can time out confusingly."""
|
||||
result = client.gateway.probe("/health/readiness/details", params=NoBody())
|
||||
result = client.proxy.probe("/health/readiness/details", params=NoBody())
|
||||
assert result.status_code == 200, (
|
||||
f"/health/readiness/details must answer 200, got {result.status_code}: {result.body[:300]}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -48,7 +48,7 @@ def _assert_otel_destination_configured(client: LoggingClient) -> None:
|
|||
"""Recorded state: the proxy reports the OTEL v2 logger among its active
|
||||
callbacks, so a missing/failed destination config fails here, before any
|
||||
traffic-based assertion can time out confusingly."""
|
||||
result = client.gateway.probe("/health/readiness/details", params=NoBody())
|
||||
result = client.proxy.probe("/health/readiness/details", params=NoBody())
|
||||
assert result.status_code == 200, (
|
||||
f"/health/readiness/details must answer 200, got {result.status_code}: {result.body[:300]}"
|
||||
)
|
||||
|
|
@ -698,13 +698,13 @@ class TestOtelTraceCompleteness:
|
|||
key = client.key_with_alias(f"otel-err-{unique_marker()}", models=[model_name])
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
|
||||
deadline = time.monotonic() + client.gateway.poll_timeout
|
||||
deadline = time.monotonic() + client.proxy.poll_timeout
|
||||
while True:
|
||||
outcome = client.chat_raw(key, model_name, "trigger an upstream auth failure", max_tokens=16)
|
||||
assert not outcome.ok, "the call must fail; the deployment's upstream key is invalid"
|
||||
if "AnthropicException" in outcome.body or time.monotonic() >= deadline:
|
||||
break
|
||||
time.sleep(client.gateway.poll_interval)
|
||||
time.sleep(client.proxy.poll_interval)
|
||||
assert "AnthropicException" in outcome.body, (
|
||||
"never saw the upstream provider failure before the deadline; the key may still be "
|
||||
f"propagating - last outcome {outcome.status_code}: {outcome.body[:200]}"
|
||||
|
|
|
|||
|
|
@ -55,13 +55,13 @@ class TestPrometheusPerKeyCardinality:
|
|||
assert response.model, f"driver call for {alias} returned no model: {response}"
|
||||
|
||||
wanted = frozenset(aliases)
|
||||
deadline = time.monotonic() + client.gateway.poll_timeout
|
||||
deadline = time.monotonic() + client.proxy.poll_timeout
|
||||
seen: frozenset[str] = frozenset()
|
||||
while time.monotonic() < deadline:
|
||||
seen = _aliases_in_metric(client.scrape_metrics(), REQUESTS_METRIC, ALIAS_LABEL)
|
||||
if wanted <= seen:
|
||||
break
|
||||
time.sleep(client.gateway.poll_interval)
|
||||
time.sleep(client.proxy.poll_interval)
|
||||
|
||||
missing = wanted - seen
|
||||
assert not missing, (
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ import pytest
|
|||
|
||||
from e2e_config import UI_BASE_URL, UI_PASSWORD, UI_USERNAME
|
||||
from management_client import ManagementClient, build_client
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from playwright.sync_api import Browser, Page
|
||||
|
|
@ -27,8 +28,8 @@ def pytest_configure(config: pytest.Config) -> None:
|
|||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def client() -> ManagementClient:
|
||||
return build_client()
|
||||
def client(proxy: ProxyClient) -> ManagementClient:
|
||||
return build_client(proxy)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
"""Client for the management-routes e2e suite: the shared Gateway plus the
|
||||
"""Client for the management-routes e2e suite: the shared ProxyClient plus the
|
||||
key/team/user/organization writes, the info/list read-backs the tests assert,
|
||||
and the raw-status calls judged by HTTP outcome (chat under a scoped key, an
|
||||
llm-only key hitting a management route).
|
||||
|
|
@ -9,7 +9,7 @@ from __future__ import annotations
|
|||
import time
|
||||
from dataclasses import dataclass
|
||||
|
||||
from e2e_gateway import Gateway, build_gateway
|
||||
from proxy_client import ProxyClient
|
||||
from e2e_http import NoBody, ProbeResult, Result, StreamingResponse, Success, UnknownApiError, unwrap
|
||||
from models import (
|
||||
ChatBody,
|
||||
|
|
@ -50,17 +50,17 @@ _TEAM_READY_SLEEP_SECONDS = 0.4
|
|||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ManagementClient:
|
||||
gateway: Gateway
|
||||
proxy: ProxyClient
|
||||
|
||||
def llm_only_key(self) -> str:
|
||||
return self.gateway.generate_key(KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"]))
|
||||
return self.proxy.generate_key(KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"]))
|
||||
|
||||
def update_key_models(self, key: str, models: list[str]) -> None:
|
||||
last: Result[NoBody] | None = None
|
||||
for attempt in range(5):
|
||||
last = self.gateway.transport.post(
|
||||
last = self.proxy.transport.post(
|
||||
"/key/update",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=KeyUpdateBody(key=key, models=models),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -79,11 +79,11 @@ class ManagementClient:
|
|||
|
||||
def delete_key_strict(self, key: str) -> None:
|
||||
"""Strict delete for the act phase of a test: a failed delete is a hard
|
||||
failure, unlike the warn-only Gateway.delete_key used at teardown."""
|
||||
failure, unlike the warn-only ProxyClient.delete_key used at teardown."""
|
||||
_ = unwrap(
|
||||
self.gateway.transport.post(
|
||||
self.proxy.transport.post(
|
||||
"/key/delete",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=KeyDeleteBody(keys=[key]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -91,9 +91,9 @@ class ManagementClient:
|
|||
|
||||
def key_alias_count(self, key_alias: str) -> int:
|
||||
return unwrap(
|
||||
self.gateway.transport.get(
|
||||
self.proxy.transport.get(
|
||||
"/key/list",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
params=KeyListParams(key_alias=key_alias),
|
||||
response_type=KeyListResponse,
|
||||
)
|
||||
|
|
@ -101,9 +101,9 @@ class ManagementClient:
|
|||
|
||||
def create_team(self, body: TeamNewBody) -> str:
|
||||
team_id = unwrap(
|
||||
self.gateway.transport.post(
|
||||
self.proxy.transport.post(
|
||||
"/team/new",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=body,
|
||||
response_type=TeamNewResponse,
|
||||
)
|
||||
|
|
@ -112,32 +112,32 @@ class ManagementClient:
|
|||
return team_id
|
||||
|
||||
def delete_team(self, team_id: str) -> None:
|
||||
_ = self.gateway.transport.post(
|
||||
_ = self.proxy.transport.post(
|
||||
"/team/delete",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=TeamDeleteBody(team_ids=[team_id]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
||||
def team_info(self, team_id: str) -> TeamData:
|
||||
return unwrap(
|
||||
self.gateway.transport.get(
|
||||
self.proxy.transport.get(
|
||||
"/team/info",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
params=TeamInfoParams(team_id=team_id),
|
||||
response_type=TeamInfoResponse,
|
||||
)
|
||||
).team_info
|
||||
|
||||
def team_info_status(self, team_id: str) -> ProbeResult:
|
||||
return self.gateway.transport.probe("/team/info", params=TeamInfoParams(team_id=team_id))
|
||||
return self.proxy.transport.probe("/team/info", params=TeamInfoParams(team_id=team_id))
|
||||
|
||||
def _wait_for_team(self, team_id: str) -> None:
|
||||
last: Result[TeamInfoResponse] | None = None
|
||||
for _ in range(_TEAM_READY_ATTEMPTS):
|
||||
last = self.gateway.transport.get(
|
||||
last = self.proxy.transport.get(
|
||||
"/team/info",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
params=TeamInfoParams(team_id=team_id),
|
||||
response_type=TeamInfoResponse,
|
||||
)
|
||||
|
|
@ -152,9 +152,9 @@ class ManagementClient:
|
|||
def add_team_member(self, team_id: str, user_id: str) -> None:
|
||||
last: Result[NoBody] | None = None
|
||||
for attempt in range(_TEAM_READY_ATTEMPTS):
|
||||
last = self.gateway.transport.post(
|
||||
last = self.proxy.transport.post(
|
||||
"/team/member_add",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=TeamMemberAddBody(team_id=team_id, member=TeamMemberEntry(role="user", user_id=user_id)),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -173,9 +173,9 @@ class ManagementClient:
|
|||
|
||||
def delete_team_member(self, team_id: str, user_id: str) -> None:
|
||||
_ = unwrap(
|
||||
self.gateway.transport.post(
|
||||
self.proxy.transport.post(
|
||||
"/team/member_delete",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=TeamMemberDeleteBody(team_id=team_id, user_id=user_id),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -183,27 +183,27 @@ class ManagementClient:
|
|||
|
||||
def create_user(self, body: UserNewBody) -> str:
|
||||
return unwrap(
|
||||
self.gateway.transport.post(
|
||||
self.proxy.transport.post(
|
||||
"/user/new",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=body,
|
||||
response_type=UserNewResponse,
|
||||
)
|
||||
).user_id
|
||||
|
||||
def delete_user(self, user_id: str) -> None:
|
||||
_ = self.gateway.transport.post(
|
||||
_ = self.proxy.transport.post(
|
||||
"/user/delete",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=UserDeleteBody(user_ids=[user_id]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
||||
def user_info(self, user_id: str) -> UserInfoResponse:
|
||||
return unwrap(
|
||||
self.gateway.transport.get(
|
||||
self.proxy.transport.get(
|
||||
"/user/info",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
params=UserInfoParams(user_id=user_id),
|
||||
response_type=UserInfoResponse,
|
||||
)
|
||||
|
|
@ -211,9 +211,9 @@ class ManagementClient:
|
|||
|
||||
def user_count(self, user_id: str) -> int:
|
||||
return unwrap(
|
||||
self.gateway.transport.get(
|
||||
self.proxy.transport.get(
|
||||
"/user/list",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
params=UserListParams(user_ids=user_id),
|
||||
response_type=UserListResponse,
|
||||
)
|
||||
|
|
@ -221,48 +221,48 @@ class ManagementClient:
|
|||
|
||||
def create_org(self, body: OrgNewBody) -> str:
|
||||
return unwrap(
|
||||
self.gateway.transport.post(
|
||||
self.proxy.transport.post(
|
||||
"/organization/new",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=body,
|
||||
response_type=OrgNewResponse,
|
||||
)
|
||||
).organization_id
|
||||
|
||||
def delete_org(self, organization_id: str) -> None:
|
||||
_ = self.gateway.transport.delete(
|
||||
_ = self.proxy.transport.delete(
|
||||
"/organization/delete",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=OrgDeleteBody(organization_ids=[organization_id]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
||||
def org_info(self, organization_id: str) -> OrgInfoResponse:
|
||||
return unwrap(
|
||||
self.gateway.transport.get(
|
||||
self.proxy.transport.get(
|
||||
"/organization/info",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
params=OrgInfoParams(organization_id=organization_id),
|
||||
response_type=OrgInfoResponse,
|
||||
)
|
||||
)
|
||||
|
||||
def chat_status(self, key: str, model: str, content: str) -> StreamingResponse:
|
||||
return self.gateway.transport.send(
|
||||
return self.proxy.transport.send(
|
||||
"/chat/completions",
|
||||
headers=self.gateway.transport.bearer(key),
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=ChatBody(model=model, messages=[ChatMessage(role="user", content=content)], max_tokens=16),
|
||||
)
|
||||
|
||||
def key_generate_status(self, key: str, body: KeyGenerateBody) -> StreamingResponse:
|
||||
return self.gateway.transport.send("/key/generate", headers=self.gateway.transport.bearer(key), json=body)
|
||||
return self.proxy.transport.send("/key/generate", headers=self.proxy.transport.bearer(key), json=body)
|
||||
|
||||
def team_new_status(self, key: str, body: TeamNewBody) -> StreamingResponse:
|
||||
return self.gateway.transport.send("/team/new", headers=self.gateway.transport.bearer(key), json=body)
|
||||
return self.proxy.transport.send("/team/new", headers=self.proxy.transport.bearer(key), json=body)
|
||||
|
||||
def user_new_status(self, key: str, body: UserNewBody) -> StreamingResponse:
|
||||
return self.gateway.transport.send("/user/new", headers=self.gateway.transport.bearer(key), json=body)
|
||||
return self.proxy.transport.send("/user/new", headers=self.proxy.transport.bearer(key), json=body)
|
||||
|
||||
|
||||
def build_client() -> ManagementClient:
|
||||
return ManagementClient(gateway=build_gateway())
|
||||
def build_client(proxy: ProxyClient) -> ManagementClient:
|
||||
return ManagementClient(proxy=proxy)
|
||||
|
|
|
|||
|
|
@ -104,8 +104,8 @@ def _provision_team(client: ManagementClient, resources: ResourceManager, alias:
|
|||
def _provision_key(
|
||||
client: ManagementClient, resources: ResourceManager, alias: str, team_id: str | None = None
|
||||
) -> str:
|
||||
key = client.gateway.generate_key(KeyGenerateBody(key_alias=alias, models=["gpt-5.5"], team_id=team_id))
|
||||
resources.defer(lambda: client.gateway.delete_key(key))
|
||||
key = client.proxy.generate_key(KeyGenerateBody(key_alias=alias, models=["gpt-5.5"], team_id=team_id))
|
||||
resources.defer(lambda: client.proxy.delete_key(key))
|
||||
return key
|
||||
|
||||
|
||||
|
|
@ -122,9 +122,9 @@ class TestKeyModelsDropdownUI:
|
|||
assert "All Team Models" not in options, f"teamless create offered 'All Team Models': {options}"
|
||||
|
||||
key = _submit_create_modal(ui_page, sentinel_label="All Proxy Models")
|
||||
resources.defer(lambda: client.gateway.delete_key(key))
|
||||
resources.defer(lambda: client.proxy.delete_key(key))
|
||||
|
||||
info = client.gateway.key_info(key)
|
||||
info = client.proxy.key_info(key)
|
||||
assert info.models == ["all-proxy-models"], f"persisted models {info.models}"
|
||||
assert info.team_id is None, f"teamless key persisted with team {info.team_id}"
|
||||
|
||||
|
|
@ -144,9 +144,9 @@ class TestKeyModelsDropdownUI:
|
|||
assert "all-proxy-models" not in options, f"team key create offered the raw sentinel: {options}"
|
||||
|
||||
key = _submit_create_modal(ui_page, sentinel_label="All Team Models")
|
||||
resources.defer(lambda: client.gateway.delete_key(key))
|
||||
resources.defer(lambda: client.proxy.delete_key(key))
|
||||
|
||||
info = client.gateway.key_info(key)
|
||||
info = client.proxy.key_info(key)
|
||||
assert info.models == ["all-team-models"], f"persisted models {info.models}"
|
||||
assert info.team_id == team_id, f"persisted team {info.team_id}, expected {team_id}"
|
||||
|
||||
|
|
|
|||
|
|
@ -27,18 +27,18 @@ from models import KeyGenerateBody, OrgNewBody, TeamNewBody, UserNewBody
|
|||
pytestmark = pytest.mark.e2e
|
||||
|
||||
def _poll[T](client: ManagementClient, attempt: Callable[[], T | None], failure: str) -> T:
|
||||
deadline = time.monotonic() + client.gateway.poll_timeout
|
||||
deadline = time.monotonic() + client.proxy.poll_timeout
|
||||
while time.monotonic() < deadline:
|
||||
found = attempt()
|
||||
if found is not None:
|
||||
return found
|
||||
time.sleep(client.gateway.poll_interval)
|
||||
time.sleep(client.proxy.poll_interval)
|
||||
pytest.fail(failure)
|
||||
|
||||
|
||||
def _generate_key(client: ManagementClient, resources: ResourceManager, body: KeyGenerateBody) -> str:
|
||||
key = client.gateway.generate_key(body)
|
||||
resources.defer(lambda: client.gateway.delete_key(key))
|
||||
key = client.proxy.generate_key(body)
|
||||
resources.defer(lambda: client.proxy.delete_key(key))
|
||||
return key
|
||||
|
||||
|
||||
|
|
@ -114,7 +114,7 @@ class TestKeyRoutes:
|
|||
KeyGenerateBody(models=["gemini-2.5-flash"], key_alias=alias, tpm_limit=424242, rpm_limit=424243),
|
||||
)
|
||||
|
||||
info = client.gateway.key_info(key)
|
||||
info = client.proxy.key_info(key)
|
||||
assert info.key_alias == alias, f"/key/info reports key_alias {info.key_alias!r}, configured {alias!r}"
|
||||
assert info.models == ["gemini-2.5-flash"], (
|
||||
f"/key/info reports models {info.models}, configured ['gemini-2.5-flash']"
|
||||
|
|
@ -143,7 +143,7 @@ class TestKeyRoutes:
|
|||
|
||||
client.update_key_models(key, ["gpt-5.5"])
|
||||
|
||||
info = client.gateway.key_info(key)
|
||||
info = client.proxy.key_info(key)
|
||||
assert info.models == ["gpt-5.5"], (
|
||||
f"/key/info reports models {info.models} after /key/update to ['gpt-5.5']"
|
||||
)
|
||||
|
|
@ -184,7 +184,7 @@ class TestTeamRoutes:
|
|||
)
|
||||
|
||||
key = _generate_key(client, resources, KeyGenerateBody(team_id=team_id))
|
||||
key_info = client.gateway.key_info(key)
|
||||
key_info = client.proxy.key_info(key)
|
||||
assert key_info.team_id == team_id, (
|
||||
f"key generated under team {team_id} carries team_id {key_info.team_id!r} in /key/info"
|
||||
)
|
||||
|
|
@ -260,7 +260,7 @@ class TestManagementRoutePermissions:
|
|||
self, client: ManagementClient, resources: ResourceManager
|
||||
) -> None:
|
||||
key = client.llm_only_key()
|
||||
resources.defer(lambda: client.gateway.delete_key(key))
|
||||
resources.defer(lambda: client.proxy.delete_key(key))
|
||||
marker = unique_marker()
|
||||
alias = f"e2e-mgmt-forbidden-key-{marker}"
|
||||
team_id = f"e2e-mgmt-forbidden-team-{marker}"
|
||||
|
|
|
|||
|
|
@ -2,15 +2,16 @@
|
|||
|
||||
The shared lifecycle (resources/scoped_key), proxy liveness handling, and the
|
||||
`e2e`/`covers` markers live in the parent tests/e2e/conftest.py. McpClient holds
|
||||
the shared Gateway, so the `resources` fixture tears down whatever this suite
|
||||
creates (keys via the Gateway, MCP servers via the deferred cleanups).
|
||||
the shared ProxyClient, so the `resources` fixture tears down whatever this suite
|
||||
creates (keys via the ProxyClient, MCP servers via the deferred cleanups).
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from mcp_client import McpClient, build_client
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def client() -> McpClient:
|
||||
return build_client()
|
||||
def client(proxy: ProxyClient) -> McpClient:
|
||||
return build_client(proxy)
|
||||
|
|
|
|||
|
|
@ -15,9 +15,9 @@ from dataclasses import dataclass
|
|||
|
||||
from pydantic import BaseModel, ConfigDict, Field, RootModel
|
||||
|
||||
from e2e_gateway import Gateway, build_gateway
|
||||
from e2e_http import Headers, NoBody, Result, unwrap
|
||||
from models import KeyGenerateBody, ObjectPermission
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
|
||||
class ApiKeyHeaders(Headers):
|
||||
|
|
@ -92,31 +92,31 @@ class McpCallToolResponse(BaseModel):
|
|||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class McpClient:
|
||||
gateway: Gateway
|
||||
proxy: ProxyClient
|
||||
|
||||
def register_server(self, *, server_name: str, alias: str, url: str) -> str:
|
||||
return unwrap(
|
||||
self.gateway.transport.post(
|
||||
self.proxy.transport.post(
|
||||
"/v1/mcp/server",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=McpServerNewBody(server_name=server_name, alias=alias, url=url),
|
||||
response_type=McpServerNewResponse,
|
||||
)
|
||||
).server_id
|
||||
|
||||
def delete_server(self, server_id: str) -> None:
|
||||
_ = self.gateway.transport.delete(
|
||||
_ = self.proxy.transport.delete(
|
||||
f"/v1/mcp/server/{server_id}",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=NoBody(),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
||||
def registered_servers(self) -> list[McpServerRow]:
|
||||
return unwrap(
|
||||
self.gateway.transport.get(
|
||||
self.proxy.transport.get(
|
||||
"/v1/mcp/server",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
params=NoBody(),
|
||||
response_type=McpServersListResponse,
|
||||
)
|
||||
|
|
@ -126,12 +126,12 @@ class McpClient:
|
|||
object_permission = (
|
||||
ObjectPermission(mcp_servers=mcp_servers) if mcp_servers is not None else None
|
||||
)
|
||||
return self.gateway.generate_key(
|
||||
return self.proxy.generate_key(
|
||||
KeyGenerateBody(models=[], user_id=user_id, object_permission=object_permission)
|
||||
)
|
||||
|
||||
def list_tools(self, key: str) -> Result[McpToolsListResponse]:
|
||||
return self.gateway.transport.get(
|
||||
return self.proxy.transport.get(
|
||||
"/mcp-rest/tools/list",
|
||||
headers=ApiKeyHeaders(x_litellm_api_key=key),
|
||||
params=NoBody(),
|
||||
|
|
@ -141,7 +141,7 @@ class McpClient:
|
|||
def call_tool(
|
||||
self, key: str, *, server_id: str, name: str, arguments: dict[str, int]
|
||||
) -> Result[McpCallToolResponse]:
|
||||
return self.gateway.transport.post(
|
||||
return self.proxy.transport.post(
|
||||
"/mcp-rest/tools/call",
|
||||
headers=ApiKeyHeaders(x_litellm_api_key=key),
|
||||
json=McpCallToolBody(name=name, arguments=arguments, server_id=server_id),
|
||||
|
|
@ -149,5 +149,5 @@ class McpClient:
|
|||
)
|
||||
|
||||
|
||||
def build_client() -> McpClient:
|
||||
return McpClient(gateway=build_gateway())
|
||||
def build_client(proxy: ProxyClient) -> McpClient:
|
||||
return McpClient(proxy=proxy)
|
||||
|
|
|
|||
|
|
@ -40,7 +40,7 @@ def _register_math_server(client: McpClient, resources: ResourceManager) -> str:
|
|||
def _key(client: McpClient, resources: ResourceManager, *, mcp_servers: list[str] | None) -> str:
|
||||
label = "allowed" if mcp_servers else "denied"
|
||||
key = client.generate_key(user_id=f"e2e-mcp-{label}-{unique_marker()}", mcp_servers=mcp_servers)
|
||||
resources.defer(lambda: client.gateway.delete_key(key))
|
||||
resources.defer(lambda: client.proxy.delete_key(key))
|
||||
return key
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -268,7 +268,7 @@ class SpendLogsParams(BaseModel):
|
|||
raise ValueError(
|
||||
"unfiltered /spend/logs returns the entire spend table and OOMs the "
|
||||
"runner on long-lived environments; filter by request_id or api_key, "
|
||||
"or use Gateway.spend_logs_window for a bounded /spend/logs/v2 read"
|
||||
"or use ProxyClient.spend_logs_window for a bounded /spend/logs/v2 read"
|
||||
)
|
||||
return self
|
||||
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
"""Gateway: the shared proxy operations, DI'd into every client (composition).
|
||||
"""ProxyClient: the shared proxy operations, DI'd into every client (composition).
|
||||
|
||||
A frozen-slots dataclass holding a Transport plus poll config. Clients hold a
|
||||
Gateway and add their own route methods; the lifecycle ResourceManager uses the
|
||||
Gateway's key/customer methods for cleanup. Read-backs are eventually consistent
|
||||
ProxyClient and add their own route methods; the lifecycle ResourceManager uses the
|
||||
ProxyClient's key/customer methods for cleanup. Read-backs are eventually consistent
|
||||
(proxy_batch_write_at ~60s) so they poll to a deadline.
|
||||
"""
|
||||
|
||||
|
|
@ -69,7 +69,7 @@ RowsPredicate = Callable[[list[SpendLogRow]], bool]
|
|||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Gateway:
|
||||
class ProxyClient:
|
||||
transport: Transport
|
||||
poll_timeout: float = 120.0
|
||||
poll_interval: float = 5.0
|
||||
|
|
@ -319,13 +319,13 @@ class Gateway:
|
|||
return self.transport.probe(path, params=params)
|
||||
|
||||
|
||||
def build_gateway(
|
||||
def build_proxy_client(
|
||||
*,
|
||||
base_url: str = PROXY_BASE_URL,
|
||||
master_key: str = MASTER_KEY,
|
||||
control_plane_base_url: str = CONTROL_PLANE_BASE_URL,
|
||||
) -> Gateway:
|
||||
"""The Gateway every suite's client is built from: a SplitTransport that routes
|
||||
) -> ProxyClient:
|
||||
"""The ProxyClient every suite's client is built from: a SplitTransport that routes
|
||||
LLM calls to the data plane (PROXY_BASE_URL) and management/admin calls to the
|
||||
control plane (CONTROL_PLANE_BASE_URL), with the shared poll budget. The two
|
||||
base URLs are the same for a monolithic proxy, so routing is then a no-op.
|
||||
|
|
@ -334,7 +334,7 @@ def build_gateway(
|
|||
way than ``e2e_config``'s env names (see ``claude_code/_env.py``); they must
|
||||
pass all three together, since a caller that overrides only the data plane
|
||||
would leave management calls pointed at the env default."""
|
||||
return Gateway(
|
||||
return ProxyClient(
|
||||
transport=SplitTransport(
|
||||
data=HttpTransport(
|
||||
base_url=base_url,
|
||||
|
|
@ -1,4 +1,4 @@
|
|||
"""Client for budget e2e tests: the shared Gateway plus budget-bearing entity
|
||||
"""Client for budget e2e tests: the shared ProxyClient plus budget-bearing entity
|
||||
management (user / team / team-member / org / customer / tag / budget-table) and
|
||||
info reads.
|
||||
|
||||
|
|
@ -15,7 +15,7 @@ from dataclasses import dataclass
|
|||
|
||||
from pydantic import AliasPath, BaseModel, Field, RootModel
|
||||
|
||||
from e2e_gateway import Gateway, build_gateway
|
||||
from proxy_client import ProxyClient
|
||||
from e2e_http import NoBody, Result, StreamingResponse, Success, unwrap
|
||||
from models import (
|
||||
AnthropicMessagesBody,
|
||||
|
|
@ -185,9 +185,9 @@ def model_budget(model: str, limit: float, period: str = "30d") -> dict[str, Mod
|
|||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BudgetClient:
|
||||
gateway: Gateway
|
||||
proxy: ProxyClient
|
||||
|
||||
# ---- generic key ops (delegate to the shared Gateway) ---------------
|
||||
# ---- generic key ops (delegate to the shared ProxyClient) ---------------
|
||||
|
||||
def generate_key(
|
||||
self,
|
||||
|
|
@ -203,7 +203,7 @@ class BudgetClient:
|
|||
budget_fallbacks: dict[str, list[str]] | None = None,
|
||||
budget_limits: list[BudgetWindow] | None = None,
|
||||
) -> str:
|
||||
return self.gateway.generate_key(
|
||||
return self.proxy.generate_key(
|
||||
KeyGenerateBody(
|
||||
models=models or [],
|
||||
max_budget=max_budget,
|
||||
|
|
@ -219,10 +219,10 @@ class BudgetClient:
|
|||
)
|
||||
|
||||
def delete_key(self, key: str) -> None:
|
||||
self.gateway.delete_key(key)
|
||||
self.proxy.delete_key(key)
|
||||
|
||||
def delete_customers(self, user_ids: list[str]) -> None:
|
||||
self.gateway.delete_customers(user_ids)
|
||||
self.proxy.delete_customers(user_ids)
|
||||
|
||||
# ---- chat (raw HTTP outcome: a budget block surfaces as a non-2xx) --
|
||||
|
||||
|
|
@ -236,9 +236,9 @@ class BudgetClient:
|
|||
user: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
) -> StreamingResponse:
|
||||
return self.gateway.transport.send(
|
||||
return self.proxy.transport.send(
|
||||
"/chat/completions",
|
||||
headers=self.gateway.transport.bearer(key),
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=ChatBody(
|
||||
model=model,
|
||||
messages=[ChatMessage(role="user", content=content)],
|
||||
|
|
@ -256,9 +256,9 @@ class BudgetClient:
|
|||
*,
|
||||
max_tokens: int = 16,
|
||||
) -> StreamingResponse:
|
||||
return self.gateway.transport.send(
|
||||
return self.proxy.transport.send(
|
||||
"/v1/messages",
|
||||
headers=self.gateway.transport.bearer(key),
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=AnthropicMessagesBody(
|
||||
model=model,
|
||||
messages=[ChatMessage(role="user", content=content)],
|
||||
|
|
@ -270,26 +270,26 @@ class BudgetClient:
|
|||
|
||||
def create_user(self, *, max_budget: float, budget_duration: str | None = None) -> str:
|
||||
return unwrap(
|
||||
self.gateway.transport.post(
|
||||
self.proxy.transport.post(
|
||||
"/user/new",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=UserNewBody(max_budget=max_budget, budget_duration=budget_duration),
|
||||
response_type=UserNewResponse,
|
||||
)
|
||||
).user_id
|
||||
|
||||
def delete_user(self, user_id: str) -> None:
|
||||
_ = self.gateway.transport.post(
|
||||
_ = self.proxy.transport.post(
|
||||
"/user/delete",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=UserDeleteBody(user_ids=[user_id]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
||||
def user_info(self, user_id: str) -> UserInfoRow | None:
|
||||
result = self.gateway.transport.get(
|
||||
result = self.proxy.transport.get(
|
||||
"/user/info",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
params=UserInfoParams(user_id=user_id),
|
||||
response_type=UserInfoResponse,
|
||||
)
|
||||
|
|
@ -302,9 +302,9 @@ class BudgetClient:
|
|||
# ---- customer / end-user -------------------------------------------
|
||||
|
||||
def create_customer(self, customer_id: str, *, max_budget: float) -> str:
|
||||
resp = self.gateway.transport.send(
|
||||
resp = self.proxy.transport.send(
|
||||
"/customer/new",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=CustomerNewBody(user_id=customer_id, max_budget=max_budget),
|
||||
)
|
||||
assert resp.ok, resp.body
|
||||
|
|
@ -314,9 +314,9 @@ class BudgetClient:
|
|||
|
||||
def create_org(self, *, max_budget: float, alias: str, budget_duration: str | None = None) -> str:
|
||||
return unwrap(
|
||||
self.gateway.transport.post(
|
||||
self.proxy.transport.post(
|
||||
"/organization/new",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=OrgNewBody(
|
||||
organization_alias=alias,
|
||||
max_budget=max_budget,
|
||||
|
|
@ -330,9 +330,9 @@ class BudgetClient:
|
|||
"""The id of the budget row backing an org; its budget_reset_at is read via
|
||||
budget_info (LIT-4570: /organization/new stores budget_duration without
|
||||
scheduling budget_reset_at, so the reset job's first tick schedules it)."""
|
||||
result = self.gateway.transport.get(
|
||||
result = self.proxy.transport.get(
|
||||
"/organization/info",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
params=OrgInfoParams(organization_id=org_id),
|
||||
response_type=OrgInfoResponse,
|
||||
)
|
||||
|
|
@ -343,9 +343,9 @@ class BudgetClient:
|
|||
return None
|
||||
|
||||
def delete_org(self, org_id: str) -> None:
|
||||
_ = self.gateway.transport.delete(
|
||||
_ = self.proxy.transport.delete(
|
||||
"/organization/delete",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=OrgDeleteBody(organization_ids=[org_id]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -362,9 +362,9 @@ class BudgetClient:
|
|||
budget_limits: list[BudgetWindow] | None = None,
|
||||
) -> str:
|
||||
team_id = unwrap(
|
||||
self.gateway.transport.post(
|
||||
self.proxy.transport.post(
|
||||
"/team/new",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=TeamNewBody(
|
||||
team_alias=alias,
|
||||
max_budget=max_budget,
|
||||
|
|
@ -379,9 +379,9 @@ class BudgetClient:
|
|||
return team_id
|
||||
|
||||
def delete_team(self, team_id: str) -> None:
|
||||
_ = self.gateway.transport.post(
|
||||
_ = self.proxy.transport.post(
|
||||
"/team/delete",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=TeamDeleteBody(team_ids=[team_id]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -389,9 +389,9 @@ class BudgetClient:
|
|||
def _wait_for_team(self, team_id: str) -> None:
|
||||
last: Result[TeamInfoResponse] | None = None
|
||||
for _ in range(_TEAM_READY_ATTEMPTS):
|
||||
last = self.gateway.transport.get(
|
||||
last = self.proxy.transport.get(
|
||||
"/team/info",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
params=TeamInfoParams(team_id=team_id),
|
||||
response_type=TeamInfoResponse,
|
||||
)
|
||||
|
|
@ -406,9 +406,9 @@ class BudgetClient:
|
|||
def add_team_member(self, team_id: str, user_id: str, *, max_budget_in_team: float | None = None) -> None:
|
||||
last_body = ""
|
||||
for attempt in range(_TEAM_READY_ATTEMPTS):
|
||||
resp = self.gateway.transport.send(
|
||||
resp = self.proxy.transport.send(
|
||||
"/team/member_add",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=TeamMemberAddBody(
|
||||
team_id=team_id,
|
||||
member=TeamMember(role="user", user_id=user_id),
|
||||
|
|
@ -432,9 +432,9 @@ class BudgetClient:
|
|||
max_budget_in_team: float | None = None,
|
||||
budget_duration: str | None = None,
|
||||
) -> None:
|
||||
resp = self.gateway.transport.send(
|
||||
resp = self.proxy.transport.send(
|
||||
"/team/member_update",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=TeamMemberUpdateBody(
|
||||
team_id=team_id,
|
||||
user_id=user_id,
|
||||
|
|
@ -448,9 +448,9 @@ class BudgetClient:
|
|||
"""The member's per-team budget_reset_at as /team/info reports it, or None if
|
||||
no reset is scheduled. The reset job advances this each time the window
|
||||
elapses; a job that skips the row leaves it pinned forever."""
|
||||
result = self.gateway.transport.get(
|
||||
result = self.proxy.transport.get(
|
||||
"/team/info",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
params=TeamInfoParams(team_id=team_id),
|
||||
response_type=TeamInfoResponse,
|
||||
)
|
||||
|
|
@ -466,18 +466,18 @@ class BudgetClient:
|
|||
# ---- tag ------------------------------------------------------------
|
||||
|
||||
def create_tag(self, name: str, *, max_budget: float) -> str:
|
||||
resp = self.gateway.transport.send(
|
||||
resp = self.proxy.transport.send(
|
||||
"/tag/new",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=TagNewBody(name=name, max_budget=max_budget),
|
||||
)
|
||||
assert resp.ok, resp.body
|
||||
return name
|
||||
|
||||
def delete_tag(self, name: str) -> None:
|
||||
_ = self.gateway.transport.post(
|
||||
_ = self.proxy.transport.post(
|
||||
"/tag/delete",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=TagDeleteBody(name=name),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -492,9 +492,9 @@ class BudgetClient:
|
|||
budget_duration: str | None = None,
|
||||
) -> str:
|
||||
return unwrap(
|
||||
self.gateway.transport.post(
|
||||
self.proxy.transport.post(
|
||||
"/budget/new",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=BudgetNewBody(
|
||||
max_budget=max_budget,
|
||||
soft_budget=soft_budget,
|
||||
|
|
@ -505,17 +505,17 @@ class BudgetClient:
|
|||
).budget_id
|
||||
|
||||
def delete_budget(self, budget_id: str) -> None:
|
||||
_ = self.gateway.transport.post(
|
||||
_ = self.proxy.transport.post(
|
||||
"/budget/delete",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=BudgetDeleteBody(id=budget_id),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
||||
def budget_info(self, budget_id: str) -> tuple[BudgetRow, ...]:
|
||||
result = self.gateway.transport.post(
|
||||
result = self.proxy.transport.post(
|
||||
"/budget/info",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=BudgetInfoBody(budgets=[budget_id]),
|
||||
response_type=BudgetInfoResponse,
|
||||
)
|
||||
|
|
@ -526,5 +526,5 @@ class BudgetClient:
|
|||
return ()
|
||||
|
||||
|
||||
def build_client() -> BudgetClient:
|
||||
return BudgetClient(gateway=build_gateway())
|
||||
def build_client(proxy: ProxyClient) -> BudgetClient:
|
||||
return BudgetClient(proxy=proxy)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
"""Budgets suite's `client` fixture.
|
||||
|
||||
The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker
|
||||
live in the parent tests/e2e/conftest.py. BudgetClient holds the shared Gateway,
|
||||
live in the parent tests/e2e/conftest.py. BudgetClient holds the shared ProxyClient,
|
||||
so the `resources` fixture cleans up keys through it; tests register entity deletes
|
||||
via `resources.defer(...)`.
|
||||
"""
|
||||
|
|
@ -9,8 +9,9 @@ via `resources.defer(...)`.
|
|||
import pytest
|
||||
|
||||
from budget_client import BudgetClient, build_client
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def client() -> BudgetClient:
|
||||
return build_client()
|
||||
def client(proxy: ProxyClient) -> BudgetClient:
|
||||
return build_client(proxy)
|
||||
|
|
|
|||
|
|
@ -30,7 +30,7 @@ def test_budget_crud_roundtrip(client: BudgetClient, resources: ResourceManager)
|
|||
# Attach the budget to a key and confirm the key reflects it.
|
||||
key = client.generate_key(budget_id=budget_id)
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
info = client.gateway.key_info(key)
|
||||
info = client.proxy.key_info(key)
|
||||
linked = info.litellm_budget_table
|
||||
assert info.budget_id == budget_id or (linked is not None and linked.max_budget == 12.5), (
|
||||
f"key does not reflect attached budget: {info.budget_id}, {linked}"
|
||||
|
|
@ -49,7 +49,7 @@ def test_budget_duration_schedules_reset_on_key(client: BudgetClient, resources:
|
|||
key = client.generate_key(max_budget=10.0, budget_duration="30d")
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
|
||||
reset_at = client.gateway.key_info(key).budget_reset_at
|
||||
reset_at = client.proxy.key_info(key).budget_reset_at
|
||||
assert reset_at, "budget_duration did not set budget_reset_at on the key"
|
||||
|
||||
# budget_duration schedules a FUTURE reset. Don't assume now+30d exactly: the
|
||||
|
|
|
|||
|
|
@ -53,7 +53,7 @@ def test_budget_fallback_reroutes_anthropic_messages_to_openai(
|
|||
|
||||
# The rerouted call must be recorded under the fallback model, not the
|
||||
# exhausted primary - proving spend tracking followed the reroute.
|
||||
rows = client.gateway.poll_logs_for_key(
|
||||
rows = client.proxy.poll_logs_for_key(
|
||||
key, predicate=lambda rows: any(FALLBACK_MODEL in (r.model or "") for r in rows)
|
||||
)
|
||||
assert any(FALLBACK_MODEL in (r.model or "") for r in rows), (
|
||||
|
|
|
|||
|
|
@ -62,7 +62,7 @@ def test_key_with_budget_duration_schedules_reset_at_creation(
|
|||
key = client.generate_key(max_budget=TINY_CAP, budget_duration=f"{WINDOW_SECONDS}s")
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
|
||||
info = client.gateway.key_info(key)
|
||||
info = client.proxy.key_info(key)
|
||||
assert info.budget_reset_at is not None, "budget_duration set no budget_reset_at"
|
||||
assert _as_datetime(info.budget_reset_at) > _as_datetime("1970-01-01T00:00:00Z")
|
||||
|
||||
|
|
@ -99,7 +99,7 @@ def test_key_budget_reset_at_advances_after_window(
|
|||
key = client.generate_key(max_budget=TINY_CAP, budget_duration=f"{WINDOW_SECONDS}s")
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
|
||||
before_raw = client.gateway.key_info(key).budget_reset_at
|
||||
before_raw = client.proxy.key_info(key).budget_reset_at
|
||||
assert before_raw is not None, "no budget_reset_at scheduled at creation"
|
||||
before = _as_datetime(before_raw)
|
||||
|
||||
|
|
@ -112,7 +112,7 @@ def test_key_budget_reset_at_advances_after_window(
|
|||
if not result.ok:
|
||||
assert is_budget_block(result), f"non-budget error during reset wait: {result.body[:200]}"
|
||||
continue
|
||||
info = client.gateway.key_info(key)
|
||||
info = client.proxy.key_info(key)
|
||||
assert info.budget_reset_at is not None, "budget_reset_at cleared by reset"
|
||||
assert _as_datetime(info.budget_reset_at) > before, (
|
||||
"budget_reset_at did not advance past the pre-reset value"
|
||||
|
|
@ -145,7 +145,7 @@ def test_multi_window_key_resets_each_window_independently(
|
|||
|
||||
start = time.monotonic()
|
||||
_drive_to_block(client, key)
|
||||
spend_at_block = client.gateway.key_info(key).spend or 0.0
|
||||
spend_at_block = client.proxy.key_info(key).spend or 0.0
|
||||
|
||||
deadline = time.monotonic() + RESET_DEADLINE_SECONDS
|
||||
while time.monotonic() < deadline:
|
||||
|
|
@ -156,7 +156,7 @@ def test_multi_window_key_resets_each_window_independently(
|
|||
assert elapsed < WINDOW_SECONDS + 90, (
|
||||
f"tight window reset took {elapsed:.0f}s - too long for {WINDOW_SECONDS}s"
|
||||
)
|
||||
assert (client.gateway.key_info(key).spend or 0.0) >= spend_at_block, (
|
||||
assert (client.proxy.key_info(key).spend or 0.0) >= spend_at_block, (
|
||||
"roomy window spend was wiped when only the tight window should reset"
|
||||
)
|
||||
return
|
||||
|
|
|
|||
|
|
@ -141,7 +141,7 @@ def test_cold_counter_reseed_keeps_counter_equal_to_db_spend(
|
|||
"the spend counter never went cold; default_redis_ttl must be short enough for it "
|
||||
"to expire, otherwise the burst reads a warm counter and the reseed is never exercised"
|
||||
)
|
||||
db_spend = client.gateway.key_info(key).spend or 0.0
|
||||
db_spend = client.proxy.key_info(key).spend or 0.0
|
||||
assert db_spend > 0, f"no DB spend accumulated from real calls: {db_spend}"
|
||||
|
||||
burst_results = []
|
||||
|
|
|
|||
|
|
@ -46,7 +46,7 @@ def member(client: BudgetClient) -> Iterator[_Member]:
|
|||
Cleanups register progressively and run LIFO best-effort through ResourceManager,
|
||||
so a partial-setup failure still releases what came before and one failed delete
|
||||
never strands the rest on the shared proxy."""
|
||||
resources = ResourceManager(client=client.gateway)
|
||||
resources = ResourceManager(client=client.proxy)
|
||||
try:
|
||||
marker = unique_marker()
|
||||
team_id = client.create_team(alias=f"e2e-team-member-{marker}", max_budget=TEAM_BUDGET)
|
||||
|
|
@ -64,7 +64,7 @@ def member(client: BudgetClient) -> Iterator[_Member]:
|
|||
def _send(client: BudgetClient, key: str) -> str | None:
|
||||
"""One member call; its response id (== the spend-log request_id) if it went
|
||||
through, else None."""
|
||||
match client.gateway.chat(
|
||||
match client.proxy.chat(
|
||||
key,
|
||||
ChatBody(
|
||||
model=MODEL,
|
||||
|
|
@ -83,7 +83,7 @@ class TestTeamMemberBudget:
|
|||
sent = frozenset(rid for rid in (_send(client, member.key) for _ in range(BURST)) if rid)
|
||||
assert sent, "no member call went through; cannot check attribution"
|
||||
|
||||
rows = client.gateway.poll_logs_for_key(
|
||||
rows = client.proxy.poll_logs_for_key(
|
||||
member.key, predicate=lambda rs: bool(sent & {r.request_id for r in rs})
|
||||
)
|
||||
logged = [row for row in rows if row.request_id in sent]
|
||||
|
|
|
|||
|
|
@ -43,7 +43,7 @@ def pair(client: BudgetClient) -> Iterator[_Pair]:
|
|||
"""One team with a large budget and two members on it: a tight member capped at
|
||||
a tiny per-team budget and a roomy member with headroom, each with their own key.
|
||||
Shared across the class and torn down LIFO best-effort when it finishes."""
|
||||
resources = ResourceManager(client=client.gateway)
|
||||
resources = ResourceManager(client=client.proxy)
|
||||
try:
|
||||
marker = unique_marker()
|
||||
team_id = client.create_team(alias=f"e2e-member-iso-{marker}", max_budget=TEAM_BUDGET)
|
||||
|
|
@ -71,7 +71,7 @@ def pair(client: BudgetClient) -> Iterator[_Pair]:
|
|||
|
||||
def _roomy_send(client: BudgetClient, key: str) -> str:
|
||||
"""One roomy-member call that must go through; returns its request id."""
|
||||
match client.gateway.chat(
|
||||
match client.proxy.chat(
|
||||
key,
|
||||
ChatBody(
|
||||
model=MODEL,
|
||||
|
|
@ -105,7 +105,7 @@ class TestTeamMemberBudgetIsolation:
|
|||
client.chat(pair.tight_key, MODEL, f"tight {unique_marker()}", max_tokens=16)
|
||||
), "tight member stopped being blocked once the peer spent"
|
||||
|
||||
rows = client.gateway.poll_logs_for_key(
|
||||
rows = client.proxy.poll_logs_for_key(
|
||||
pair.roomy_key, predicate=lambda rs: bool(sent & {r.request_id for r in rs})
|
||||
)
|
||||
logged = [row for row in rows if row.request_id in sent]
|
||||
|
|
|
|||
|
|
@ -1,15 +1,16 @@
|
|||
"""Quota-management suite's `client` fixture.
|
||||
|
||||
The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker
|
||||
live in the parent tests/e2e/conftest.py. QuotaClient holds the shared Gateway,
|
||||
live in the parent tests/e2e/conftest.py. QuotaClient holds the shared ProxyClient,
|
||||
so the `resources` fixture cleans up keys through it.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from quota_client import QuotaClient, build_client
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def client() -> QuotaClient:
|
||||
return build_client()
|
||||
def client(proxy: ProxyClient) -> QuotaClient:
|
||||
return build_client(proxy)
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
"""Client for the quota-management suite: the shared Gateway plus raw chat
|
||||
"""Client for the quota-management suite: the shared ProxyClient plus raw chat
|
||||
calls judged by HTTP status, body, and headers (a rate-limit block is a 429
|
||||
whose body and retry-after header carry the contract, not a typed success
|
||||
model)."""
|
||||
|
|
@ -7,19 +7,19 @@ from __future__ import annotations
|
|||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from e2e_gateway import Gateway, build_gateway
|
||||
from proxy_client import ProxyClient
|
||||
from e2e_http import StreamingResponse
|
||||
from models import ChatBody, ChatMessage
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class QuotaClient:
|
||||
gateway: Gateway
|
||||
proxy: ProxyClient
|
||||
|
||||
def chat(self, key: str, model: str, content: str, *, max_tokens: int = 16) -> StreamingResponse:
|
||||
return self.gateway.transport.send(
|
||||
return self.proxy.transport.send(
|
||||
"/chat/completions",
|
||||
headers=self.gateway.transport.bearer(key),
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=ChatBody(
|
||||
model=model,
|
||||
messages=[ChatMessage(role="user", content=content)],
|
||||
|
|
@ -28,5 +28,5 @@ class QuotaClient:
|
|||
)
|
||||
|
||||
|
||||
def build_client() -> QuotaClient:
|
||||
return QuotaClient(gateway=build_gateway())
|
||||
def build_client(proxy: ProxyClient) -> QuotaClient:
|
||||
return QuotaClient(proxy=proxy)
|
||||
|
|
|
|||
|
|
@ -131,8 +131,8 @@ def _limited_key(
|
|||
rpm_limit: int | None = None,
|
||||
tpm_limit: int | None = None,
|
||||
) -> str:
|
||||
key = client.gateway.generate_key(KeyGenerateBody(models=[MODEL], rpm_limit=rpm_limit, tpm_limit=tpm_limit))
|
||||
resources.defer(lambda: client.gateway.delete_key(key))
|
||||
key = client.proxy.generate_key(KeyGenerateBody(models=[MODEL], rpm_limit=rpm_limit, tpm_limit=tpm_limit))
|
||||
resources.defer(lambda: client.proxy.delete_key(key))
|
||||
return key
|
||||
|
||||
|
||||
|
|
@ -147,7 +147,7 @@ def _first_ok(client: QuotaClient, key: str) -> _FirstOk:
|
|||
cache picks it up, so retry on 401 to a deadline; a 401 never reaches the
|
||||
rate limiter, so only the successful call consumes budget. Any other failure
|
||||
is behavior under test and fails hard."""
|
||||
deadline = time.monotonic() + client.gateway.poll_timeout
|
||||
deadline = time.monotonic() + client.proxy.poll_timeout
|
||||
while True:
|
||||
sent_at = time.monotonic()
|
||||
outcome = _chat(client, key)
|
||||
|
|
@ -155,7 +155,7 @@ def _first_ok(client: QuotaClient, key: str) -> _FirstOk:
|
|||
return _FirstOk(sent_at=sent_at, response=outcome)
|
||||
if outcome.status_code != 401 or time.monotonic() >= deadline:
|
||||
require_successful_call(outcome)
|
||||
time.sleep(client.gateway.poll_interval)
|
||||
time.sleep(client.proxy.poll_interval)
|
||||
|
||||
|
||||
def _assert_rate_limited(outcome: StreamingResponse, limit_type: str) -> None:
|
||||
|
|
@ -178,7 +178,7 @@ class TestKeyRateLimits:
|
|||
@pytest.mark.covers("quota_management.ratelimit.rpm.blocks_over_limit")
|
||||
def test_rpm_limit_blocks_over_limit(self, client: QuotaClient, resources: ResourceManager) -> None:
|
||||
key = _limited_key(client, resources, rpm_limit=3)
|
||||
info = client.gateway.key_info(key)
|
||||
info = client.proxy.key_info(key)
|
||||
assert info.rpm_limit == 3, f"/key/info reports rpm_limit {info.rpm_limit}, configured 3"
|
||||
|
||||
_ = _first_ok(client, key)
|
||||
|
|
@ -190,7 +190,7 @@ class TestKeyRateLimits:
|
|||
@pytest.mark.covers("quota_management.ratelimit.tpm.blocks_over_limit")
|
||||
def test_tpm_limit_blocks_over_limit(self, client: QuotaClient, resources: ResourceManager) -> None:
|
||||
key = _limited_key(client, resources, tpm_limit=TPM_LIMIT)
|
||||
info = client.gateway.key_info(key)
|
||||
info = client.proxy.key_info(key)
|
||||
assert info.tpm_limit == TPM_LIMIT, f"/key/info reports tpm_limit {info.tpm_limit}, configured {TPM_LIMIT}"
|
||||
|
||||
first = _first_ok(client, key)
|
||||
|
|
@ -213,7 +213,7 @@ class TestKeyRateLimits:
|
|||
first = _first_ok(client, key)
|
||||
_assert_rate_limited(_chat(client, key), "requests")
|
||||
|
||||
deadline = time.monotonic() + client.gateway.poll_timeout
|
||||
deadline = time.monotonic() + client.proxy.poll_timeout
|
||||
while time.monotonic() < deadline:
|
||||
attempt_sent_at = time.monotonic()
|
||||
outcome = _chat(client, key)
|
||||
|
|
@ -228,7 +228,7 @@ class TestKeyRateLimits:
|
|||
assert outcome.status_code == 429, (
|
||||
f"while the window drains only 429s are acceptable, got {outcome.status_code}: {outcome.body[:300]}"
|
||||
)
|
||||
time.sleep(client.gateway.poll_interval)
|
||||
time.sleep(client.proxy.poll_interval)
|
||||
pytest.fail("a blocked key never recovered after the rate-limit window elapsed")
|
||||
|
||||
@pytest.mark.covers("quota_management.ratelimit.rpm.headers_report_remaining")
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
"""Spend-tracking suite's `client` fixture and driver-model registration.
|
||||
|
||||
The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker
|
||||
live in the parent tests/e2e/conftest.py. SpendClient exposes the shared Gateway
|
||||
(GatewayProvider), so the `resources` fixture cleans up keys and customers this
|
||||
live in the parent tests/e2e/conftest.py. SpendClient exposes the shared ProxyClient
|
||||
(ProxyClientProvider), so the `resources` fixture cleans up keys and customers this
|
||||
suite creates.
|
||||
|
||||
The suite drives real calls through three deployments. On the stage gateway they
|
||||
|
|
@ -21,6 +21,7 @@ import pytest
|
|||
|
||||
from models import LiteLLMParamsBody
|
||||
from spend_e2e_client import SpendClient, build_client
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
|
||||
def _driver_params(provider_model: str, env_var: str) -> LiteLLMParamsBody:
|
||||
|
|
@ -38,18 +39,18 @@ DRIVER_MODELS: tuple[tuple[str, str, str], ...] = (
|
|||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def client() -> SpendClient:
|
||||
return build_client()
|
||||
def client(proxy: ProxyClient) -> SpendClient:
|
||||
return build_client(proxy)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session", autouse=True)
|
||||
def driver_models(client: SpendClient) -> Iterator[None]:
|
||||
existing = frozenset(entry.model_name for entry in client.gateway.model_info())
|
||||
existing = frozenset(entry.model_name for entry in client.proxy.model_info())
|
||||
created = tuple(
|
||||
client.gateway.create_model(name, _driver_params(provider_model, env_var))
|
||||
client.proxy.create_model(name, _driver_params(provider_model, env_var))
|
||||
for name, provider_model, env_var in DRIVER_MODELS
|
||||
if name not in existing
|
||||
)
|
||||
yield
|
||||
for model_id in created:
|
||||
client.gateway.delete_model(model_id)
|
||||
client.proxy.delete_model(model_id)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
"""Spend-tracking e2e client: a Gateway plus the spend-specific read endpoints.
|
||||
"""Spend-tracking e2e client: a ProxyClient plus the spend-specific read endpoints.
|
||||
|
||||
Generic proxy operations (keys, customers, chat/embed, route probing, SpendLogs
|
||||
polling) come from the shared Gateway, DI'd in (composition, not inheritance).
|
||||
polling) come from the shared ProxyClient, DI'd in (composition, not inheritance).
|
||||
This client adds only the spend surface: /spend/calculate, /spend/tags,
|
||||
key-spend polling, and the route probes the breadth test uses.
|
||||
|
||||
|
|
@ -27,7 +27,7 @@ from e2e_http import (
|
|||
is_ok,
|
||||
unwrap,
|
||||
)
|
||||
from e2e_gateway import Gateway, build_gateway
|
||||
from proxy_client import ProxyClient
|
||||
from models import (
|
||||
ChatBody,
|
||||
ChatMessage,
|
||||
|
|
@ -96,7 +96,7 @@ def _chat_body(
|
|||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SpendClient:
|
||||
gateway: Gateway
|
||||
proxy: ProxyClient
|
||||
|
||||
def chat(
|
||||
self,
|
||||
|
|
@ -108,19 +108,19 @@ class SpendClient:
|
|||
tags: list[str] | None = None,
|
||||
user: str | None = None,
|
||||
) -> Result[ChatResponse]:
|
||||
return self.gateway.chat(
|
||||
return self.proxy.chat(
|
||||
key, _chat_body(model, content, max_tokens=max_tokens, tags=tags, user=user)
|
||||
)
|
||||
|
||||
def chat_stream(
|
||||
self, key: str, model: str, content: str, *, max_tokens: int | None = None
|
||||
) -> StreamingResponse:
|
||||
return self.gateway.chat_stream(
|
||||
return self.proxy.chat_stream(
|
||||
key, _chat_body(model, content, max_tokens=max_tokens, stream=True)
|
||||
)
|
||||
|
||||
def embed(self, key: str, model: str, content: str) -> Result[EmbedResponse]:
|
||||
return self.gateway.embed(key, EmbedBody(model=model, input=content))
|
||||
return self.proxy.embed(key, EmbedBody(model=model, input=content))
|
||||
|
||||
def poll_logs_for_key(
|
||||
self,
|
||||
|
|
@ -129,15 +129,15 @@ class SpendClient:
|
|||
min_rows: int = 1,
|
||||
predicate: Callable[[list[SpendLogRow]], bool] | None = None,
|
||||
) -> list[SpendLogRow]:
|
||||
return self.gateway.poll_logs_for_key(
|
||||
return self.proxy.poll_logs_for_key(
|
||||
key, min_rows=min_rows, predicate=predicate
|
||||
)
|
||||
|
||||
def calculate_spend(self, model: str, content: str) -> float:
|
||||
return unwrap(
|
||||
self.gateway.transport.post(
|
||||
self.proxy.transport.post(
|
||||
"/spend/calculate",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
json=SpendCalculateBody(
|
||||
model=model, messages=[ChatMessage(role="user", content=content)]
|
||||
),
|
||||
|
|
@ -146,9 +146,9 @@ class SpendClient:
|
|||
).cost
|
||||
|
||||
def spend_by_tags(self) -> list[TagSpend]:
|
||||
result = self.gateway.transport.get(
|
||||
result = self.proxy.transport.get(
|
||||
"/spend/tags",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
params=NoBody(),
|
||||
response_type=SpendTagsResponse,
|
||||
)
|
||||
|
|
@ -160,7 +160,7 @@ class SpendClient:
|
|||
|
||||
def poll_tag_spend(self, tag: str, *, minimum: float = 0.0) -> TagSpend | None:
|
||||
"""Poll /spend/tags until the tag's aggregate reaches `minimum`; last seen."""
|
||||
deadline = time.monotonic() + self.gateway.poll_timeout
|
||||
deadline = time.monotonic() + self.proxy.poll_timeout
|
||||
entry: TagSpend | None = None
|
||||
while time.monotonic() < deadline:
|
||||
matches = [
|
||||
|
|
@ -170,17 +170,17 @@ class SpendClient:
|
|||
entry = matches[0]
|
||||
if (entry.total_spend or 0.0) >= minimum:
|
||||
return entry
|
||||
time.sleep(self.gateway.poll_interval)
|
||||
time.sleep(self.proxy.poll_interval)
|
||||
return entry
|
||||
|
||||
def poll_key_spend(self, key: str, *, minimum: float = 0.0) -> float:
|
||||
deadline = time.monotonic() + self.gateway.poll_timeout
|
||||
deadline = time.monotonic() + self.proxy.poll_timeout
|
||||
spend = 0.0
|
||||
while time.monotonic() < deadline:
|
||||
spend = self.gateway.key_info(key).spend or 0.0
|
||||
spend = self.proxy.key_info(key).spend or 0.0
|
||||
if spend > minimum:
|
||||
return spend
|
||||
time.sleep(self.gateway.poll_interval)
|
||||
time.sleep(self.proxy.poll_interval)
|
||||
return spend
|
||||
|
||||
def spend_logs_page(
|
||||
|
|
@ -191,9 +191,9 @@ class SpendClient:
|
|||
now = datetime.now(timezone.utc)
|
||||
fmt = "%Y-%m-%d %H:%M:%S"
|
||||
return unwrap(
|
||||
self.gateway.transport.get(
|
||||
self.proxy.transport.get(
|
||||
"/spend/logs/v2",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
params=SpendLogsPageParams(
|
||||
start_date=(now - timedelta(days=1)).strftime(fmt),
|
||||
end_date=(now + timedelta(days=1)).strftime(fmt),
|
||||
|
|
@ -206,18 +206,18 @@ class SpendClient:
|
|||
)
|
||||
|
||||
def probe(self, path: str, *, params: DateRangeParams) -> ProbeResult:
|
||||
return self.gateway.transport.probe(path, params=params)
|
||||
return self.proxy.transport.probe(path, params=params)
|
||||
|
||||
def openapi(self) -> OpenAPISchema:
|
||||
return unwrap(
|
||||
self.gateway.transport.get(
|
||||
self.proxy.transport.get(
|
||||
"/openapi.json",
|
||||
headers=self.gateway.transport.master,
|
||||
headers=self.proxy.transport.master,
|
||||
params=NoBody(),
|
||||
response_type=OpenAPISchema,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def build_client() -> SpendClient:
|
||||
return SpendClient(gateway=build_gateway())
|
||||
def build_client(proxy: ProxyClient) -> SpendClient:
|
||||
return SpendClient(proxy=proxy)
|
||||
|
|
|
|||
|
|
@ -475,12 +475,12 @@ def test_spend_logs_endpoint_returns_spend(
|
|||
)
|
||||
)
|
||||
|
||||
gateway = client.gateway
|
||||
deadline = time.monotonic() + gateway.poll_timeout
|
||||
proxy = client.proxy
|
||||
deadline = time.monotonic() + proxy.poll_timeout
|
||||
while True:
|
||||
result = gateway.transport.get(
|
||||
result = proxy.transport.get(
|
||||
"/spend/logs",
|
||||
headers=gateway.transport.master,
|
||||
headers=proxy.transport.master,
|
||||
params=SpendLogsParams(api_key=scoped_key),
|
||||
response_type=SpendLogs,
|
||||
)
|
||||
|
|
@ -493,4 +493,4 @@ def test_spend_logs_endpoint_returns_spend(
|
|||
f"/spend/logs never surfaced the key's spend before the deadline; "
|
||||
f"saw {_summarize(rows)}"
|
||||
)
|
||||
time.sleep(gateway.poll_interval)
|
||||
time.sleep(proxy.poll_interval)
|
||||
|
|
|
|||
|
|
@ -1,20 +1,20 @@
|
|||
"""Client for the complexity auto-router e2e tests.
|
||||
|
||||
The suite drives the shared /chat/completions and spend-log reads on the Gateway,
|
||||
so this client only carries the Gateway the shared lifecycle needs for cleanup.
|
||||
The suite drives the shared /chat/completions and spend-log reads on the ProxyClient,
|
||||
so this client only carries the ProxyClient the shared lifecycle needs for cleanup.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from e2e_gateway import Gateway, build_gateway
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ComplexityRouterClient:
|
||||
gateway: Gateway
|
||||
proxy: ProxyClient
|
||||
|
||||
|
||||
def build_client() -> ComplexityRouterClient:
|
||||
return ComplexityRouterClient(gateway=build_gateway())
|
||||
def build_client(proxy: ProxyClient) -> ComplexityRouterClient:
|
||||
return ComplexityRouterClient(proxy=proxy)
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker
|
||||
live in the parent tests/e2e/conftest.py. ComplexityRouterClient holds the shared
|
||||
Gateway, so the `resources` fixture cleans up keys this suite creates.
|
||||
ProxyClient, so the `resources` fixture cleans up keys this suite creates.
|
||||
|
||||
Also registers `complexity-smart-router` via management /model/new when the
|
||||
proxy does not already list it (compose has it in static config; stage does not).
|
||||
|
|
@ -16,7 +16,7 @@ import pytest
|
|||
from requests import RequestException
|
||||
|
||||
from complexity_router_client import ComplexityRouterClient, build_client
|
||||
from e2e_gateway import Gateway
|
||||
from proxy_client import ProxyClient
|
||||
from e2e_http import NoBody, Success
|
||||
from lifecycle import ResourceManager
|
||||
from models import (
|
||||
|
|
@ -46,27 +46,27 @@ ROUTER_KEY_MODELS = [ROUTER_MODEL, "gpt-5.5", "claude-haiku-4-5"]
|
|||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def client() -> ComplexityRouterClient:
|
||||
return build_client()
|
||||
def client(proxy: ProxyClient) -> ComplexityRouterClient:
|
||||
return build_client(proxy)
|
||||
|
||||
|
||||
def _model_is_servable(gateway: Gateway, model_name: str) -> bool:
|
||||
result = gateway.transport.get(
|
||||
def _model_is_servable(proxy: ProxyClient, model_name: str) -> bool:
|
||||
result = proxy.transport.get(
|
||||
"/v1/models",
|
||||
headers=gateway.transport.master,
|
||||
headers=proxy.transport.master,
|
||||
params=NoBody(),
|
||||
response_type=ModelsListResponse,
|
||||
)
|
||||
return isinstance(result, Success) and any(entry.id == model_name for entry in result.data.data)
|
||||
|
||||
|
||||
def _router_is_callable(gateway: Gateway) -> bool:
|
||||
def _router_is_callable(proxy: ProxyClient) -> bool:
|
||||
"""True only when a short chat against the virtual router succeeds; every error
|
||||
(the Invalid-model-name reload race, but also 401, 5xx, and network) counts as
|
||||
not-callable so infra/auth blips can't be mistaken for a working router."""
|
||||
key = gateway.generate_key(KeyGenerateBody(models=ROUTER_KEY_MODELS, user_id="e2e-complexity-probe"))
|
||||
key = proxy.generate_key(KeyGenerateBody(models=ROUTER_KEY_MODELS, user_id="e2e-complexity-probe"))
|
||||
try:
|
||||
result = gateway.chat(
|
||||
result = proxy.chat(
|
||||
key,
|
||||
ChatBody(
|
||||
model=ROUTER_MODEL,
|
||||
|
|
@ -75,7 +75,7 @@ def _router_is_callable(gateway: Gateway) -> bool:
|
|||
),
|
||||
)
|
||||
finally:
|
||||
gateway.delete_key(key)
|
||||
proxy.delete_key(key)
|
||||
return isinstance(result, Success)
|
||||
|
||||
|
||||
|
|
@ -86,18 +86,18 @@ def _ensure_complexity_smart_router( # pyright: ignore[reportUnusedFunction] #
|
|||
"""Ensure the complexity router virtual model exists for this session.
|
||||
|
||||
Compose already declares it in docker-compose.yml; stage does not. Register
|
||||
via Gateway.create_model (waits for data-plane /v1/models) when missing, then
|
||||
via ProxyClient.create_model (waits for data-plane /v1/models) when missing, then
|
||||
probe a real chat so a list-only false positive cannot pass the fixture.
|
||||
"""
|
||||
gateway = client.gateway
|
||||
if _model_is_servable(gateway, ROUTER_MODEL) and _router_is_callable(gateway):
|
||||
proxy = client.proxy
|
||||
if _model_is_servable(proxy, ROUTER_MODEL) and _router_is_callable(proxy):
|
||||
yield
|
||||
return
|
||||
|
||||
try:
|
||||
model_id = gateway.create_model(ROUTER_MODEL, ROUTER_PARAMS)
|
||||
model_id = proxy.create_model(ROUTER_MODEL, ROUTER_PARAMS)
|
||||
except (AssertionError, RequestException) as exc:
|
||||
if _model_is_servable(gateway, ROUTER_MODEL) and _router_is_callable(gateway):
|
||||
if _model_is_servable(proxy, ROUTER_MODEL) and _router_is_callable(proxy):
|
||||
yield
|
||||
return
|
||||
raise AssertionError(
|
||||
|
|
@ -106,7 +106,7 @@ def _ensure_complexity_smart_router( # pyright: ignore[reportUnusedFunction] #
|
|||
) from exc
|
||||
|
||||
try:
|
||||
if not _router_is_callable(gateway):
|
||||
if not _router_is_callable(proxy):
|
||||
raise AssertionError(
|
||||
f"{ROUTER_MODEL!r} registered as {model_id!r} and listed on "
|
||||
f"/v1/models but chat still returns Invalid model name; "
|
||||
|
|
@ -114,14 +114,14 @@ def _ensure_complexity_smart_router( # pyright: ignore[reportUnusedFunction] #
|
|||
)
|
||||
yield
|
||||
finally:
|
||||
gateway.delete_model(model_id)
|
||||
proxy.delete_model(model_id)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def complexity_key(resources: ResourceManager, client: ComplexityRouterClient) -> str:
|
||||
"""Per-test key allowed to call the complexity router and its tier backends."""
|
||||
key = client.gateway.generate_key(
|
||||
key = client.proxy.generate_key(
|
||||
KeyGenerateBody(models=ROUTER_KEY_MODELS, user_id="e2e-complexity-router")
|
||||
)
|
||||
resources.defer(lambda: client.gateway.delete_key(key))
|
||||
resources.defer(lambda: client.proxy.delete_key(key))
|
||||
return key
|
||||
|
|
|
|||
|
|
@ -48,7 +48,7 @@ class TestComplexityRouterLlmClassifier:
|
|||
self, client: ComplexityRouterClient, complexity_key: str
|
||||
) -> None:
|
||||
chat = unwrap(
|
||||
client.gateway.chat(
|
||||
client.proxy.chat(
|
||||
complexity_key,
|
||||
ChatBody(
|
||||
model=ROUTER_MODEL,
|
||||
|
|
@ -59,7 +59,7 @@ class TestComplexityRouterLlmClassifier:
|
|||
)
|
||||
assert chat.choices, f"router returned no choices: {chat}"
|
||||
|
||||
rows = client.gateway.poll_logs_for_key(complexity_key, min_rows=1)
|
||||
rows = client.proxy.poll_logs_for_key(complexity_key, min_rows=1)
|
||||
served = [row.model for row in rows]
|
||||
# Exactly one spend row for the routed completion (not the classifier sub-call).
|
||||
# Membership allows alias vs provider-prefixed forms across compose and stage.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue