From c4f9688c07010062bc63358da4423c5d8e0233ae Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 31 Jul 2026 21:14:26 -0700 Subject: [PATCH] perf(e2e/claude_code): run the compat matrix concurrently end to end The matrix had three serial stretches left. `run_compat.sh` only ever handed six of the sixteen feature directories to `pytest -n`, so the other ten columns ran one cell at a time; the ten HTTP-probe cells looped their three Claude tiers in sequence, paying the per-provider rate limiter's wait three times per cell; and the session-scoped fixture that registers the fifteen compat deployments POSTed them one at a time, in every xdist worker process Registration moves to `pytest_configure` on the controller and goes out in parallel, so a run registers each deployment once instead of once per worker. Workers learn the controller owns registration through `pytest_configure_node`; a run whose controller never loaded this conftest still registers per worker, so no invocation loses its deployments. The probe cells share one `run_probe_cell` body that fans the tiers out and reports rows in declared order, replacing ten copies of the same loop --- tests/e2e/CLAUDE.md | 2 +- tests/e2e/claude_code/_compat_models.py | 87 +++++++- tests/e2e/claude_code/_probe_cell.py | 100 +++++++++ tests/e2e/claude_code/conftest.py | 192 ++++++++++++------ .../count_tokens/test_anthropic.py | 30 ++- .../claude_code/count_tokens/test_azure.py | 30 ++- .../count_tokens/test_bedrock_converse.py | 30 ++- .../count_tokens/test_bedrock_invoke.py | 30 ++- .../count_tokens/test_vertex_ai.py | 30 ++- tests/e2e/claude_code/run_compat.sh | 28 +-- tests/e2e/claude_code/test_compat_models.py | 135 ++++++++++++ tests/e2e/claude_code/test_probe_cell.py | 166 +++++++++++++++ .../claude_code/tool_search/test_anthropic.py | 34 ++-- .../e2e/claude_code/tool_search/test_azure.py | 34 ++-- .../tool_search/test_bedrock_converse.py | 34 ++-- .../tool_search/test_bedrock_invoke.py | 34 ++-- .../claude_code/tool_search/test_vertex_ai.py | 34 ++-- 17 files changed, 768 insertions(+), 262 deletions(-) create mode 100644 tests/e2e/claude_code/_probe_cell.py create mode 100644 tests/e2e/claude_code/test_compat_models.py create mode 100644 tests/e2e/claude_code/test_probe_cell.py diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 180639b53e3..ff4e3deb501 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -21,7 +21,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `load/` - throughput/performance under concurrency: drives real concurrent traffic through the whole stack with Locust and asserts a throughput SLO; marked `load` so the parent conftest collects it last and it never perturbs latency-sensitive suites. Also home of the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`): Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; additionally marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, because it spends real provider money (driven by `.github/workflows/weekly_load_anomaly.yml`) - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests -- `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke +- `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher, and `run_compat.sh` runs the whole matrix under `pytest -n` (the cross-process rate limiter in `rate_limiter.py` bounds the aggregate per-provider request rate however wide the fanout goes). The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` ## MCP suite: real Datadog only diff --git a/tests/e2e/claude_code/_compat_models.py b/tests/e2e/claude_code/_compat_models.py index 23c9c82e596..2e7ea1e7816 100644 --- a/tests/e2e/claude_code/_compat_models.py +++ b/tests/e2e/claude_code/_compat_models.py @@ -1,5 +1,5 @@ -"""Load the claude_code compat matrix's deployment list from -``test_config.yaml``. +"""Load and register the claude_code compat matrix's deployment list +from ``test_config.yaml``. ``test_config.yaml`` is the ground-truth config the stage deployment uses; parsing it at fixture time means a change there (new tier, tier @@ -12,9 +12,10 @@ collection instead of at 400-time. from __future__ import annotations +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from pathlib import Path -from typing import Callable, Mapping +from typing import Callable, Mapping, Sequence import yaml @@ -74,6 +75,86 @@ def load_all_deployments( ) +@dataclass(frozen=True, slots=True) +class DeploymentFailure: + model_name: str + reason: str + + +@dataclass(frozen=True, slots=True) +class RegistrationOutcome: + model_ids: tuple[str, ...] + failures: tuple[DeploymentFailure, ...] + + +RegisterDeployment = Callable[[CompatDeployment], str] +DeleteModel = Callable[[str], None] + + +def register_deployments( + deployments: Sequence[CompatDeployment], + register: RegisterDeployment, +) -> RegistrationOutcome: + """Register every deployment concurrently and collect the outcome. + + Registration is one ``/model/new`` POST plus a poll until the data + plane lists the model, so fifteen of them in sequence cost fifteen + round trips before the first cell can run. They are independent, so + they go out at once and the run pays roughly one. + + Errors stay values: a deployment the proxy can't serve (missing + provider credential, bad params) lands in ``failures`` instead of + aborting the batch, because the cells that target it should fail + loudly on their own rather than taking the whole matrix down. + """ + if not deployments: + return RegistrationOutcome(model_ids=(), failures=()) + + def _one(deployment: CompatDeployment) -> str | DeploymentFailure: + try: + return register(deployment) + except Exception as exc: + return DeploymentFailure( + model_name=deployment.model_name, + reason=f"{type(exc).__name__}: {exc}", + ) + + with ThreadPoolExecutor(max_workers=len(deployments)) as pool: + outcomes = tuple(pool.map(_one, deployments)) + + return RegistrationOutcome( + model_ids=tuple(o for o in outcomes if isinstance(o, str)), + failures=tuple(o for o in outcomes if isinstance(o, DeploymentFailure)), + ) + + +def unregister_deployments( + model_ids: Sequence[str], + delete: DeleteModel, +) -> tuple[DeploymentFailure, ...]: + """Delete every registered deployment concurrently; return what failed. + + Teardown runs after the last cell, so its cost is pure wall time at + the end of a run. Failures come back as values because one flaky + delete must not mask the test results the run just produced. + """ + if not model_ids: + return () + + def _one(model_id: str) -> DeploymentFailure | None: + try: + delete(model_id) + except Exception as exc: + return DeploymentFailure( + model_name=model_id, + reason=f"{type(exc).__name__}: {exc}", + ) + return None + + with ThreadPoolExecutor(max_workers=len(model_ids)) as pool: + return tuple(o for o in pool.map(_one, model_ids) if o is not None) + + def all_expected_model_names( *, config_path: Path = CONFIG_PATH, diff --git a/tests/e2e/claude_code/_probe_cell.py b/tests/e2e/claude_code/_probe_cell.py new file mode 100644 index 00000000000..aeebe92a340 --- /dev/null +++ b/tests/e2e/claude_code/_probe_cell.py @@ -0,0 +1,100 @@ +"""Shared body for the HTTP-probe compat cells (`count_tokens`, `tool_search`). + +Every probe row does the same three things: fire one request per Claude +tier, shape-check each response, and report one `compat_result` row per +tier so the matrix builder's "all tiers must pass" aggregator still sees +three rows for the cell. + +The tiers are independent HTTP round trips, so they run concurrently +here for the same reason the CLI rows fan out in +`run_claude_models_parallel`: a cell's wall time should be one probe, +not the sum of three. That matters more than the raw request latency +suggests, because every probe first blocks on the cross-process +per-provider token bucket (`rate_limiter.py`) -- a serial cell paid +three of those waits back to back while holding a pytest worker idle. + +Report order follows the declared model order, not completion order, so +the results artifact is deterministic no matter which tier answers +first. + +The leading underscore in the filename is what keeps pytest from +collecting this module as a test file. +""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from typing import Callable, Sequence + +import pytest +from pydantic import BaseModel + +from e2e_http import NetworkError, Result + + +def probe_models_parallel[R: BaseModel]( + *, + models: Sequence[str], + probe: Callable[[str], Result[R]], +) -> dict[str, Result[R]]: + """Run `probe` once per model concurrently and return outcomes keyed by model. + + Errors stay values: a probe that raises (the rate limiter does file + I/O, and `infer_provider` can reject an edge-case model string) + becomes a `NetworkError` entry rather than an exception that would + discard the other tiers' outcomes. + """ + if not models: + raise ValueError("models must be a non-empty sequence") + + def _one(model: str) -> tuple[str, Result[R]]: + try: + return model, probe(model) + except Exception as exc: + return model, NetworkError( + message=f"unexpected error probing {model!r}: {type(exc).__name__}: {exc}" + ) + + with ThreadPoolExecutor(max_workers=len(models)) as pool: + return dict(pool.map(_one, models)) + + +def _failure_message(model: str, probe_name: str, error: str) -> str: + return f"[{model}] {probe_name} probe failed: {error}" + + +def run_probe_cell[R: BaseModel]( + *, + compat_result, + models: Sequence[str], + probe: Callable[[str], Result[R]], + check_shape: Callable[[Result[R]], str | None], + probe_name: str, +) -> None: + """Run the shared HTTP-probe cell body across every tier in `models`. + + `probe` and `check_shape` are the per-feature halves: the caller + binds the proxy client and api key into `probe`, and passes the + matching `assert_*_shape` as `check_shape`. Both are plain callables + so a test can inject fakes without a live proxy. + """ + outcomes = probe_models_parallel(models=models, probe=probe) + checked = tuple((model, check_shape(outcomes[model])) for model in models) + + for model, error in checked: + compat_result.add( + {"status": "pass"} + if error is None + else { + "status": "fail", + "error": _failure_message(model, probe_name, error), + } + ) + + failures = tuple( + _failure_message(model, probe_name, error) + for model, error in checked + if error is not None + ) + if failures: + pytest.fail("; ".join(failures), pytrace=False) diff --git a/tests/e2e/claude_code/conftest.py b/tests/e2e/claude_code/conftest.py index aaad7667936..603adc70106 100644 --- a/tests/e2e/claude_code/conftest.py +++ b/tests/e2e/claude_code/conftest.py @@ -1,6 +1,6 @@ """Pytest plumbing for the Claude Code compatibility matrix. -Three responsibilities live here: +Four responsibilities live here: 1. The `compat_result` fixture — the only API a test author needs to learn. Tests call `compat_result.set({"status": "pass"})` (or fail / not_applicable) @@ -22,6 +22,12 @@ Three responsibilities live here: summary that the binary-search helper consumes to decide whether the current X/Y/Z values were too aggressive. +4. Compat model registration — `pytest_configure` POSTs every deployment + in `test_config.yaml` to the proxy (in parallel) and + `pytest_unconfigure` deletes them again. Doing it on the controller + rather than in a session fixture keeps it to one registration pass + per run instead of one per xdist worker. + The (feature, provider) inference comes from the test file path: the parent directory name is the feature_id (matching `manifest.yaml`), and the file stem after the leading `test_` is the provider id. This avoids per-file @@ -298,19 +304,19 @@ def pytest_runtest_makereport(item, call): ) -def _is_xdist_worker(session) -> bool: - """Return True iff the current pytest session is an xdist worker. +def _is_xdist_worker(config) -> bool: + """Return True iff this pytest process is an xdist worker. The standard idiom is to look up `workerinput` on the config; the controller process doesn't have it, the workers do. We deliberately don't `import xdist` because the suite must keep running when xdist isn't installed at all. """ - return hasattr(session.config, "workerinput") + return hasattr(config, "workerinput") -def _xdist_worker_id(session) -> Optional[str]: - info = getattr(session.config, "workerinput", None) +def _xdist_worker_id(config) -> Optional[str]: + info = getattr(config, "workerinput", None) if not info: return None return info.get("workerid") @@ -452,7 +458,7 @@ def pytest_sessionstart(session): _COLLECTOR.items.clear() _manifest_feature_ids.cache_clear() - if _is_xdist_worker(session): + if _is_xdist_worker(session.config): return artifact_path = Path(os.environ.get(RESULTS_ARTIFACT_ENV) or DEFAULT_ARTIFACT_PATH) shard_dir = _shard_dir(artifact_path) @@ -494,11 +500,11 @@ def pytest_sessionfinish(session, exitstatus): artifact_path = Path(os.environ.get(RESULTS_ARTIFACT_ENV) or DEFAULT_ARTIFACT_PATH) shard_dir = _shard_dir(artifact_path) has_worker_shards = shard_dir.is_dir() and any(shard_dir.glob("*.json")) - if not _COLLECTOR.items and not _is_xdist_worker(session) and not has_worker_shards: + if not _COLLECTOR.items and not _is_xdist_worker(session.config) and not has_worker_shards: return shard_dir.mkdir(parents=True, exist_ok=True) - worker_id = _xdist_worker_id(session) or "main" + worker_id = _xdist_worker_id(session.config) or "main" shard_path = shard_dir / f"{worker_id}.json" shard_path.write_text( json.dumps( @@ -514,7 +520,7 @@ def pytest_sessionfinish(session, exitstatus): # Workers stop here. The controller merges; if we're not running # under xdist, we are effectively the controller. - if _is_xdist_worker(session): + if _is_xdist_worker(session.config): return merged_rows: List[Dict[str, Any]] = [] @@ -551,31 +557,44 @@ def pytest_sessionfinish(session, exitstatus): # --------------------------------------------------------------------------- -# Session-scoped compat model registration. +# Run-scoped compat model registration. # # The compat cells probe hardcoded virtual names like ``claude-sonnet-4-5`` # and ``claude-sonnet-4-5-bedrock-invoke``. On stage those live in the # gateway's model_list at deploy time; locally the docker-config.yaml # under tests/e2e/ only declares one of them, so every non-haiku cell -# 400s with ``Invalid model name``. The fixture here reconciles the two: -# it reads ``test_config.yaml`` (the ground-truth compat matrix config) -# and POSTs ``/model/new`` for the subset whose provider credentials are -# actually set in the current environment, then tears them all down at -# session end. +# 400s with ``Invalid model name``. The hooks here reconcile the two: +# they read ``test_config.yaml`` (the ground-truth compat matrix config), +# POST ``/model/new`` for every deployment it declares, and delete them +# again when the run ends. +# +# This runs in ``pytest_configure`` rather than a session fixture so it +# happens once per run instead of once per process: under ``-n``, a +# session fixture executes in every worker, so a 12-worker matrix run +# registered (and later deleted) the same 15 deployments 12 times over. +# The controller configures before it forks any worker, so registering +# there also means the deployments are already servable by the time the +# first cell runs. The 15 POSTs go out concurrently, so startup costs +# roughly one round trip rather than fifteen. # # Kept below the rest of the conftest so the compat-artifact hooks stay -# grouped up top. The fixture is opt-in via autouse=True on the session -# scope, so a cell that hits the proxy sees the deployment ready without -# any per-cell wiring, and pure unit tests that never reach the proxy -# pay only one skipped-liveness check. +# grouped up top. # --------------------------------------------------------------------------- +from typing import TYPE_CHECKING # noqa: E402 + from claude_code._env import ProxyConfig, resolve_proxy # noqa: E402 from claude_code._compat_models import ( # noqa: E402 CompatDeployment, + RegistrationOutcome, load_all_deployments, + register_deployments, + unregister_deployments, ) +if TYPE_CHECKING: + from proxy_client import ProxyClient + def _build_control_plane_client(proxy_config: ProxyConfig): """Local import of the shared harness so the pure-unit-test tree @@ -606,14 +625,58 @@ def _register_deployment(proxy, deployment: CompatDeployment) -> str: ) -@pytest.fixture(scope="session", autouse=True) -def _compat_models_registered() -> Any: - """Register every compat deployment against the running proxy, then - tear them all down on session exit. +@dataclass(frozen=True, slots=True) +class _RegisteredCompatModels: + proxy: "ProxyClient" + model_ids: Tuple[str, ...] - Skips silently if the proxy env is not configured (no - ``LITELLM_PROXY_URL``/``LITELLM_MASTER_KEY``) so unit-test runs - stay hermetic. + +_REGISTERED_MODELS = pytest.StashKey[_RegisteredCompatModels]() +_CONTROLLER_REGISTERED = "claude_code_compat_models_registered" + + +@pytest.hookimpl(optionalhook=True) +def pytest_configure_node(node) -> None: + """Tell each xdist worker the controller owns registration. + + xdist-only hook (`optionalhook` so the suite still loads without + xdist installed). The controller only reaches it when this conftest + was loaded before the workers forked, which is exactly when its + `pytest_configure` registered the deployments - so the flag doubles + as the workers' "don't register your own copies" signal. A run whose + controller never loaded this conftest (e.g. `pytest tests/e2e -n + auto`, where the directory is only reached during each worker's own + collection) never sets it, and the workers fall back to registering + for themselves.""" + node.workerinput[_CONTROLLER_REGISTERED] = True + + +def _controller_already_registered(config) -> bool: + workerinput = getattr(config, "workerinput", None) + return bool(workerinput) and bool(workerinput.get(_CONTROLLER_REGISTERED)) + + +def _report_registration_failures(outcome: RegistrationOutcome) -> None: + summary = "\n".join( + f" - {failure.model_name}: {failure.reason}" for failure in outcome.failures + ) + total = len(outcome.failures) + len(outcome.model_ids) + print( + f"[compat] {len(outcome.failures)} of {total} deployments failed to " + f"register (proxy likely missing that provider's credentials); cells " + f"that target them will fail loudly:\n{summary}", + file=sys.stderr, + flush=True, + ) + + +def pytest_configure(config: pytest.Config) -> None: + """Register every compat deployment against the running proxy, once per run. + + Does nothing if the proxy env is not configured (no + ``LITELLM_PROXY_URL``/``LITELLM_MASTER_KEY``) so unit-test runs stay + hermetic, in a ``--collect-only`` run, which executes no cell, or in + an xdist worker whose controller already registered. Design note: we always attempt to register all 15 deployments, regardless of what credentials are exported in the test-runner's @@ -621,44 +684,51 @@ def _compat_models_registered() -> Any: (via docker-compose ``env_file``), not the shell running pytest - so gating on shell env would filter out deployments the proxy can actually serve. Per-deployment ``/model/new`` failures are printed - but do not abort the session: the cells that need that specific + but do not abort the run: 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).""" + if _controller_already_registered(config) or config.getoption( + "collectonly", default=False + ): + return proxy_config = resolve_proxy() if proxy_config is None: - yield return - from requests import RequestException - 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(proxy, deployment) - registered_ids.append(model_id) - except (AssertionError, RequestException) as exc: - failures.append((deployment.model_name, str(exc))) - if failures: - summary = "\n".join( - f" - {name}: {reason}" for name, reason in failures - ) - print( - f"[compat fixture] {len(failures)} of " - f"{len(failures) + len(registered_ids)} deployments " - f"failed to register (proxy likely missing that provider's " - f"credentials); cells that target them will fail loudly:\n" - f"{summary}" - ) - yield - finally: - for model_id in registered_ids: - try: - proxy.delete_model(model_id) - except (AssertionError, RequestException): - # Best-effort — teardown surfaces via warnings inside - # ``delete_model`` already; swallowing here so one flaky - # delete does not mask real test failures. - pass + outcome = register_deployments( + load_all_deployments(), + lambda deployment: _register_deployment(proxy, deployment), + ) + config.stash[_REGISTERED_MODELS] = _RegisteredCompatModels( + proxy=proxy, + model_ids=outcome.model_ids, + ) + if outcome.failures: + _report_registration_failures(outcome) + + +def pytest_unconfigure(config: pytest.Config) -> None: + """Delete the deployments this process registered in `pytest_configure`. + + On the controller that lands after every worker has finished, so no + cell can still be routing to a deployment we are tearing down. + Delete failures are reported, never raised - a flaky teardown must + not mask the results the run just produced.""" + registered = config.stash.get(_REGISTERED_MODELS, None) + if registered is None: + return + failures = unregister_deployments( + registered.model_ids, + registered.proxy.delete_model, + ) + if failures: + summary = "\n".join( + f" - {failure.model_name}: {failure.reason}" for failure in failures + ) + print( + f"[compat] {len(failures)} deployment(s) failed to delete; they " + f"will linger on the proxy:\n{summary}", + file=sys.stderr, + flush=True, + ) diff --git a/tests/e2e/claude_code/count_tokens/test_anthropic.py b/tests/e2e/claude_code/count_tokens/test_anthropic.py index 05110d24e86..18d613ad052 100644 --- a/tests/e2e/claude_code/count_tokens/test_anthropic.py +++ b/tests/e2e/claude_code/count_tokens/test_anthropic.py @@ -30,9 +30,9 @@ diff. The cell goes red if *any* tier's probe fails the minimal shape check; the matrix's per-cell aggregator handles that automatically. -Three tiers run sequentially because count_tokens is cheap (<100ms -per request typical) and the parallelization that matters for the -CLI rows isn't useful here. +The three tiers are probed concurrently via `run_probe_cell`: each +probe first waits on the shared per-provider token bucket, so a +sequential cell paid that wait three times before reporting. """ from __future__ import annotations @@ -40,6 +40,7 @@ from __future__ import annotations import pytest from claude_code._env import require_proxy_client +from claude_code._probe_cell import run_probe_cell from claude_code.http_probe import ( assert_count_tokens_shape, probe_count_tokens, @@ -59,19 +60,12 @@ def test_count_tokens_anthropic(compat_result): assert the response shape.""" client, api_key = require_proxy_client(compat_result) - failures = [] - for model in ANTHROPIC_MODELS: - result = probe_count_tokens( + run_probe_cell( + compat_result=compat_result, + models=ANTHROPIC_MODELS, + probe=lambda model: probe_count_tokens( client=client, api_key=api_key, model=model - ) - shape_error = assert_count_tokens_shape(result) - if shape_error is not None: - error = f"[{model}] count_tokens probe failed: {shape_error}" - compat_result.add({"status": "fail", "error": error}) - failures.append(error) - continue - - compat_result.add({"status": "pass"}) - - if failures: - pytest.fail("; ".join(failures), pytrace=False) + ), + check_shape=assert_count_tokens_shape, + probe_name="count_tokens", + ) diff --git a/tests/e2e/claude_code/count_tokens/test_azure.py b/tests/e2e/claude_code/count_tokens/test_azure.py index c60c623ae89..f4c3bb83e92 100644 --- a/tests/e2e/claude_code/count_tokens/test_azure.py +++ b/tests/e2e/claude_code/count_tokens/test_azure.py @@ -30,9 +30,9 @@ diff. The cell goes red if *any* tier's probe fails the minimal shape check; the matrix's per-cell aggregator handles that automatically. -Three tiers run sequentially because count_tokens is cheap (<100ms -per request typical) and the parallelization that matters for the -CLI rows isn't useful here. +The three tiers are probed concurrently via `run_probe_cell`: each +probe first waits on the shared per-provider token bucket, so a +sequential cell paid that wait three times before reporting. """ from __future__ import annotations @@ -40,6 +40,7 @@ from __future__ import annotations import pytest from claude_code._env import require_proxy_client +from claude_code._probe_cell import run_probe_cell from claude_code.http_probe import ( assert_count_tokens_shape, probe_count_tokens, @@ -59,19 +60,12 @@ def test_count_tokens_azure(compat_result): assert the response shape.""" client, api_key = require_proxy_client(compat_result) - failures = [] - for model in AZURE_MODELS: - result = probe_count_tokens( + run_probe_cell( + compat_result=compat_result, + models=AZURE_MODELS, + probe=lambda model: probe_count_tokens( client=client, api_key=api_key, model=model - ) - shape_error = assert_count_tokens_shape(result) - if shape_error is not None: - error = f"[{model}] count_tokens probe failed: {shape_error}" - compat_result.add({"status": "fail", "error": error}) - failures.append(error) - continue - - compat_result.add({"status": "pass"}) - - if failures: - pytest.fail("; ".join(failures), pytrace=False) + ), + check_shape=assert_count_tokens_shape, + probe_name="count_tokens", + ) diff --git a/tests/e2e/claude_code/count_tokens/test_bedrock_converse.py b/tests/e2e/claude_code/count_tokens/test_bedrock_converse.py index 0cb4766bb31..cec544150d1 100644 --- a/tests/e2e/claude_code/count_tokens/test_bedrock_converse.py +++ b/tests/e2e/claude_code/count_tokens/test_bedrock_converse.py @@ -30,9 +30,9 @@ diff. The cell goes red if *any* tier's probe fails the minimal shape check; the matrix's per-cell aggregator handles that automatically. -Three tiers run sequentially because count_tokens is cheap (<100ms -per request typical) and the parallelization that matters for the -CLI rows isn't useful here. +The three tiers are probed concurrently via `run_probe_cell`: each +probe first waits on the shared per-provider token bucket, so a +sequential cell paid that wait three times before reporting. """ from __future__ import annotations @@ -40,6 +40,7 @@ from __future__ import annotations import pytest from claude_code._env import require_proxy_client +from claude_code._probe_cell import run_probe_cell from claude_code.http_probe import ( assert_count_tokens_shape, probe_count_tokens, @@ -59,19 +60,12 @@ def test_count_tokens_bedrock_converse(compat_result): assert the response shape.""" client, api_key = require_proxy_client(compat_result) - failures = [] - for model in BEDROCK_CONVERSE_MODELS: - result = probe_count_tokens( + run_probe_cell( + compat_result=compat_result, + models=BEDROCK_CONVERSE_MODELS, + probe=lambda model: probe_count_tokens( client=client, api_key=api_key, model=model - ) - shape_error = assert_count_tokens_shape(result) - if shape_error is not None: - error = f"[{model}] count_tokens probe failed: {shape_error}" - compat_result.add({"status": "fail", "error": error}) - failures.append(error) - continue - - compat_result.add({"status": "pass"}) - - if failures: - pytest.fail("; ".join(failures), pytrace=False) + ), + check_shape=assert_count_tokens_shape, + probe_name="count_tokens", + ) diff --git a/tests/e2e/claude_code/count_tokens/test_bedrock_invoke.py b/tests/e2e/claude_code/count_tokens/test_bedrock_invoke.py index f1389574527..a3353b48e47 100644 --- a/tests/e2e/claude_code/count_tokens/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/count_tokens/test_bedrock_invoke.py @@ -30,9 +30,9 @@ diff. The cell goes red if *any* tier's probe fails the minimal shape check; the matrix's per-cell aggregator handles that automatically. -Three tiers run sequentially because count_tokens is cheap (<100ms -per request typical) and the parallelization that matters for the -CLI rows isn't useful here. +The three tiers are probed concurrently via `run_probe_cell`: each +probe first waits on the shared per-provider token bucket, so a +sequential cell paid that wait three times before reporting. """ from __future__ import annotations @@ -40,6 +40,7 @@ from __future__ import annotations import pytest from claude_code._env import require_proxy_client +from claude_code._probe_cell import run_probe_cell from claude_code.http_probe import ( assert_count_tokens_shape, probe_count_tokens, @@ -59,19 +60,12 @@ def test_count_tokens_bedrock_invoke(compat_result): assert the response shape.""" client, api_key = require_proxy_client(compat_result) - failures = [] - for model in BEDROCK_INVOKE_MODELS: - result = probe_count_tokens( + run_probe_cell( + compat_result=compat_result, + models=BEDROCK_INVOKE_MODELS, + probe=lambda model: probe_count_tokens( client=client, api_key=api_key, model=model - ) - shape_error = assert_count_tokens_shape(result) - if shape_error is not None: - error = f"[{model}] count_tokens probe failed: {shape_error}" - compat_result.add({"status": "fail", "error": error}) - failures.append(error) - continue - - compat_result.add({"status": "pass"}) - - if failures: - pytest.fail("; ".join(failures), pytrace=False) + ), + check_shape=assert_count_tokens_shape, + probe_name="count_tokens", + ) diff --git a/tests/e2e/claude_code/count_tokens/test_vertex_ai.py b/tests/e2e/claude_code/count_tokens/test_vertex_ai.py index 0894214d4f0..462ac9d63c0 100644 --- a/tests/e2e/claude_code/count_tokens/test_vertex_ai.py +++ b/tests/e2e/claude_code/count_tokens/test_vertex_ai.py @@ -30,9 +30,9 @@ diff. The cell goes red if *any* tier's probe fails the minimal shape check; the matrix's per-cell aggregator handles that automatically. -Three tiers run sequentially because count_tokens is cheap (<100ms -per request typical) and the parallelization that matters for the -CLI rows isn't useful here. +The three tiers are probed concurrently via `run_probe_cell`: each +probe first waits on the shared per-provider token bucket, so a +sequential cell paid that wait three times before reporting. """ from __future__ import annotations @@ -40,6 +40,7 @@ from __future__ import annotations import pytest from claude_code._env import require_proxy_client +from claude_code._probe_cell import run_probe_cell from claude_code.http_probe import ( assert_count_tokens_shape, probe_count_tokens, @@ -60,19 +61,12 @@ def test_count_tokens_vertex_ai(compat_result): assert the response shape.""" client, api_key = require_proxy_client(compat_result) - failures = [] - for model in VERTEX_AI_MODELS: - result = probe_count_tokens( + run_probe_cell( + compat_result=compat_result, + models=VERTEX_AI_MODELS, + probe=lambda model: probe_count_tokens( client=client, api_key=api_key, model=model - ) - shape_error = assert_count_tokens_shape(result) - if shape_error is not None: - error = f"[{model}] count_tokens probe failed: {shape_error}" - compat_result.add({"status": "fail", "error": error}) - failures.append(error) - continue - - compat_result.add({"status": "pass"}) - - if failures: - pytest.fail("; ".join(failures), pytrace=False) + ), + check_shape=assert_count_tokens_shape, + probe_name="count_tokens", + ) diff --git a/tests/e2e/claude_code/run_compat.sh b/tests/e2e/claude_code/run_compat.sh index b881cf1d31e..4386bc62479 100755 --- a/tests/e2e/claude_code/run_compat.sh +++ b/tests/e2e/claude_code/run_compat.sh @@ -34,6 +34,10 @@ # Optional env (parallelism): # COMPAT_XDIST_WORKERS passed to `pytest -n` (default: auto) # +# Optional env (scope): +# COMPAT_TARGETS space-separated pytest targets +# (default: the whole matrix) +# # Optional env (artifacts): # COMPAT_RESULTS_PATH default: compat-results.json # COMPAT_RATE_LIMIT_SUMMARY_PATH default: compat-rate-limit-summary.json @@ -58,6 +62,12 @@ fi # this is a "go as fast as the limiter allows" knob, not a tuning knob. workers="${COMPAT_XDIST_WORKERS:-auto}" +# What to run. The whole suite by default: every feature directory is a +# column of the published matrix, and running a subset silently leaves +# those cells "not_tested". Override with COMPAT_TARGETS to iterate on +# one feature (e.g. COMPAT_TARGETS=tests/e2e/claude_code/thinking). +read -r -a targets <<< "${COMPAT_TARGETS:-tests/e2e/claude_code}" + # Where the artifacts land. We resolve them now so the summary file is # always at a known path the caller can grep, even if they didn't set # the env explicitly. @@ -71,28 +81,22 @@ for provider in ANTHROPIC AZURE VERTEX_AI BEDROCK_CONVERSE BEDROCK_INVOKE OPENAI done echo " BURST=${LITELLM_COMPAT_RATE_BURST:-default(=rate)}" echo "[run_compat] xdist workers: ${workers}" +echo "[run_compat] targets: ${targets[*]}" echo "[run_compat] results: ${results_path}" echo "[run_compat] summary: ${summary_path}" -# Run only the per-feature live tests; skip the unit-test directories -# (they're under directories starting with `_`). The dist=loadfile -# scheduler keeps each test file pinned to a single worker, which is -# what we want — every test in a file shares a single ThreadPoolExecutor -# fanout, and we don't gain anything by splitting it across workers. +# Every cell is one test, so `dist=load` hands each one to whichever +# worker is free next; the rate limiter keeps the aggregate per-provider +# request rate bounded no matter how many run at once. start=$(date +%s) set +e COMPAT_RESULTS_PATH="${results_path}" \ COMPAT_RATE_LIMIT_SUMMARY_PATH="${summary_path}" \ PATH="$HOME/.local/bin:$PATH" \ uv run pytest \ - tests/e2e/claude_code/basic_messaging_non_streaming \ - tests/e2e/claude_code/basic_messaging_streaming \ - tests/e2e/claude_code/thinking \ - tests/e2e/claude_code/tool_use \ - tests/e2e/claude_code/vision \ - tests/e2e/claude_code/prompt_caching_5m \ + "${targets[@]}" \ -n "${workers}" \ - --dist=loadfile \ + --dist=load \ -q \ "$@" diff --git a/tests/e2e/claude_code/test_compat_models.py b/tests/e2e/claude_code/test_compat_models.py new file mode 100644 index 00000000000..cb7c255e524 --- /dev/null +++ b/tests/e2e/claude_code/test_compat_models.py @@ -0,0 +1,135 @@ +"""Harness tests for compat deployment registration. + +No `e2e` marker and no proxy: `register_deployments` / +`unregister_deployments` take the "register one" and "delete one" +callables as parameters, so the concurrency and the errors-as-values +contract are pinned here with fakes. +""" + +from __future__ import annotations + +import threading +from typing import List + +import pytest + +from claude_code._compat_models import ( + CompatDeployment, + DeploymentFailure, + RegistrationOutcome, + register_deployments, + unregister_deployments, +) +from models import LiteLLMParamsBody + + +BARRIER_TIMEOUT_SECONDS = 10.0 + + +def _deployments(count: int) -> tuple[CompatDeployment, ...]: + return tuple( + CompatDeployment( + model_name=f"model-{index}", + litellm_params=LiteLLMParamsBody(model=f"anthropic/model-{index}"), + ) + for index in range(count) + ) + + +def test_deployments_register_concurrently() -> None: + """Registration is one POST plus a poll per deployment, so they all go + out at once: the barrier only trips if every registration is in + flight together.""" + deployments = _deployments(4) + barrier = threading.Barrier(len(deployments), timeout=BARRIER_TIMEOUT_SECONDS) + + def register(deployment: CompatDeployment) -> str: + barrier.wait() + return f"id-{deployment.model_name}" + + outcome = register_deployments(deployments, register) + + assert outcome.failures == () + assert sorted(outcome.model_ids) == [f"id-model-{index}" for index in range(4)] + + +def test_one_unservable_deployment_does_not_abort_the_batch() -> None: + """A deployment the proxy has no credential for must come back as a + failure value; the cells that target it fail loudly on their own, + and every other deployment still registers.""" + deployments = _deployments(3) + + def register(deployment: CompatDeployment) -> str: + if deployment.model_name == "model-1": + raise AssertionError("no credential on the proxy") + return f"id-{deployment.model_name}" + + outcome = register_deployments(deployments, register) + + assert outcome.failures == ( + DeploymentFailure( + model_name="model-1", + reason="AssertionError: no credential on the proxy", + ), + ) + assert sorted(outcome.model_ids) == ["id-model-0", "id-model-2"] + + +def test_registering_nothing_touches_the_proxy_not_at_all() -> None: + def register(deployment: CompatDeployment) -> str: + raise AssertionError(f"must not register {deployment.model_name}") + + assert register_deployments((), register) == RegistrationOutcome( + model_ids=(), failures=() + ) + + +def test_deletes_run_concurrently() -> None: + model_ids = ("id-1", "id-2", "id-3") + barrier = threading.Barrier(len(model_ids), timeout=BARRIER_TIMEOUT_SECONDS) + + def delete(model_id: str) -> None: + barrier.wait() + + assert unregister_deployments(model_ids, delete) == () + + +def test_a_failed_delete_is_reported_and_the_rest_still_run() -> None: + deleted: List[str] = [] + lock = threading.Lock() + + def delete(model_id: str) -> None: + with lock: + deleted.append(model_id) + if model_id == "id-2": + raise RuntimeError("already gone") + + failures = unregister_deployments(("id-1", "id-2", "id-3"), delete) + + assert sorted(deleted) == ["id-1", "id-2", "id-3"] + assert failures == ( + DeploymentFailure(model_name="id-2", reason="RuntimeError: already gone"), + ) + + +def test_deleting_nothing_is_a_no_op() -> None: + def delete(model_id: str) -> None: + raise AssertionError(f"must not delete {model_id}") + + assert unregister_deployments((), delete) == () + + +@pytest.mark.parametrize("count", [1, 5]) +def test_every_deployment_is_registered_exactly_once(count: int) -> None: + registered: List[str] = [] + lock = threading.Lock() + + def register(deployment: CompatDeployment) -> str: + with lock: + registered.append(deployment.model_name) + return f"id-{deployment.model_name}" + + outcome = register_deployments(_deployments(count), register) + + assert sorted(registered) == [f"model-{index}" for index in range(count)] + assert len(outcome.model_ids) == count diff --git a/tests/e2e/claude_code/test_probe_cell.py b/tests/e2e/claude_code/test_probe_cell.py new file mode 100644 index 00000000000..25026a36c89 --- /dev/null +++ b/tests/e2e/claude_code/test_probe_cell.py @@ -0,0 +1,166 @@ +"""Harness tests for the shared HTTP-probe cell body. + +No `e2e` marker and no proxy: these drive `_probe_cell` with injected +fakes to pin the three properties the cells depend on - the tiers are +probed concurrently, rows are reported in declared order rather than +completion order, and a probe that raises becomes a value instead of +losing the other tiers' outcomes. Same shape of harness coverage as +`coverage_registry/test_collector.py`. +""" + +from __future__ import annotations + +import threading +import time +from typing import Any, Callable, Dict, List + +import pytest +from pydantic import BaseModel + +from claude_code._probe_cell import probe_models_parallel, run_probe_cell +from e2e_http import NetworkError, Result, Success + + +MODELS = ["claude-haiku-4-5", "claude-sonnet-4-5", "claude-opus-4-7"] +BARRIER_TIMEOUT_SECONDS = 10.0 + + +class _Body(BaseModel): + model: str + + +def _ok(model: str) -> Result[_Body]: + return Success(status_code=200, data=_Body(model=model)) + + +def _probed_model(result: Result[_Body]) -> str | None: + match result: + case Success(data=data): + return data.model + case _: + return None + + +def _network_error(result: Result[_Body]) -> str | None: + match result: + case NetworkError(message=message): + return message + case _: + return None + + +class _Recorder: + """Stand-in for the `compat_result` fixture: keeps the rows a cell reports.""" + + def __init__(self) -> None: + self.rows: List[Dict[str, Any]] = [] + + def add(self, row: Dict[str, Any]) -> None: + self.rows.append(row) + + +def _staggered_probe(delays: Dict[str, float]) -> Callable[[str], Result[_Body]]: + """A probe whose completion order is the reverse of the declared order.""" + + def probe(model: str) -> Result[_Body]: + time.sleep(delays[model]) + return _ok(model) + + return probe + + +def test_tiers_are_probed_concurrently() -> None: + """Every tier must be in flight at once: the barrier only trips if all + three probes are running together, so a sequential implementation + times out on the first wait.""" + barrier = threading.Barrier(len(MODELS), timeout=BARRIER_TIMEOUT_SECONDS) + + def probe(model: str) -> Result[_Body]: + barrier.wait() + return _ok(model) + + outcomes = probe_models_parallel(models=MODELS, probe=probe) + + assert [_probed_model(outcomes[model]) for model in MODELS] == MODELS + + +def test_rows_follow_declared_order_not_completion_order() -> None: + """The results artifact must be deterministic, so rows are reported in + the order the cell declares its tiers even when the slowest tier is + declared first.""" + recorder = _Recorder() + delays = { + "claude-haiku-4-5": 0.15, + "claude-sonnet-4-5": 0.05, + "claude-opus-4-7": 0.0, + } + + def check_shape(result: Result[_Body]) -> str | None: + model = _probed_model(result) + return None if model == "claude-sonnet-4-5" else f"bad {model}" + + with pytest.raises(pytest.fail.Exception) as excinfo: + run_probe_cell( + compat_result=recorder, + models=MODELS, + probe=_staggered_probe(delays), + check_shape=check_shape, + probe_name="count_tokens", + ) + + assert recorder.rows == [ + { + "status": "fail", + "error": ( + "[claude-haiku-4-5] count_tokens probe failed: " + "bad claude-haiku-4-5" + ), + }, + {"status": "pass"}, + { + "status": "fail", + "error": ( + "[claude-opus-4-7] count_tokens probe failed: bad claude-opus-4-7" + ), + }, + ] + assert str(excinfo.value).index("claude-haiku-4-5") < str(excinfo.value).index( + "claude-opus-4-7" + ) + + +def test_all_tiers_passing_reports_one_pass_row_each() -> None: + recorder = _Recorder() + + run_probe_cell( + compat_result=recorder, + models=MODELS, + probe=_ok, + check_shape=lambda _: None, + probe_name="tool_search", + ) + + assert recorder.rows == [{"status": "pass"}] * len(MODELS) + + +def test_a_raising_probe_becomes_a_value_and_spares_the_other_tiers() -> None: + """One tier blowing up (rate-limiter file I/O, an unmappable model id) + must not discard the tiers that answered.""" + + def probe(model: str) -> Result[_Body]: + if model == "claude-sonnet-4-5": + raise RuntimeError("boom") + return _ok(model) + + outcomes = probe_models_parallel(models=MODELS, probe=probe) + + assert _probed_model(outcomes["claude-haiku-4-5"]) == "claude-haiku-4-5" + assert _probed_model(outcomes["claude-opus-4-7"]) == "claude-opus-4-7" + message = _network_error(outcomes["claude-sonnet-4-5"]) + assert message is not None + assert "RuntimeError: boom" in message + + +def test_no_models_is_a_programming_error() -> None: + with pytest.raises(ValueError): + probe_models_parallel(models=[], probe=_ok) diff --git a/tests/e2e/claude_code/tool_search/test_anthropic.py b/tests/e2e/claude_code/tool_search/test_anthropic.py index a23d1b3d2bb..99c4c3483cd 100644 --- a/tests/e2e/claude_code/tool_search/test_anthropic.py +++ b/tests/e2e/claude_code/tool_search/test_anthropic.py @@ -34,11 +34,11 @@ and forwarding, and the upstream either accepts or 400s. A red cell here is always a proxy-side regression, not a flaky model-behavior artifact. -Three Claude tiers are probed in sequence (count is too low to be -worth the parallelism overhead, and HTTP probes don't compete for -the proxy's `--num-workers` slots the way CLI subprocess runs do). -The matrix's "all three must pass" rule still applies via the -per-cell aggregator. +The three Claude tiers are probed concurrently via `run_probe_cell`, +the same way the CLI rows fan out; each probe first waits on the +shared per-provider token bucket, so a sequential cell paid that wait +three times before reporting. The matrix's "all three must pass" rule +still applies via the per-cell aggregator. """ from __future__ import annotations @@ -46,6 +46,7 @@ from __future__ import annotations import pytest from claude_code._env import require_proxy_client +from claude_code._probe_cell import run_probe_cell from claude_code.http_probe import ( assert_tool_search_shape, probe_tool_search, @@ -66,17 +67,12 @@ def test_tool_search_anthropic(compat_result): tier.""" client, api_key = require_proxy_client(compat_result) - failures = [] - for model in ANTHROPIC_MODELS: - result = probe_tool_search(client=client, api_key=api_key, model=model) - shape_error = assert_tool_search_shape(result) - if shape_error is not None: - error = f"[{model}] tool_search probe failed: {shape_error}" - compat_result.add({"status": "fail", "error": error}) - failures.append(error) - continue - - compat_result.add({"status": "pass"}) - - if failures: - pytest.fail("; ".join(failures), pytrace=False) + run_probe_cell( + compat_result=compat_result, + models=ANTHROPIC_MODELS, + probe=lambda model: probe_tool_search( + client=client, api_key=api_key, model=model + ), + check_shape=assert_tool_search_shape, + probe_name="tool_search", + ) diff --git a/tests/e2e/claude_code/tool_search/test_azure.py b/tests/e2e/claude_code/tool_search/test_azure.py index b094a35ea63..8fc6db12ade 100644 --- a/tests/e2e/claude_code/tool_search/test_azure.py +++ b/tests/e2e/claude_code/tool_search/test_azure.py @@ -34,11 +34,11 @@ and forwarding, and the upstream either accepts or 400s. A red cell here is always a proxy-side regression, not a flaky model-behavior artifact. -Three Claude tiers are probed in sequence (count is too low to be -worth the parallelism overhead, and HTTP probes don't compete for -the proxy's `--num-workers` slots the way CLI subprocess runs do). -The matrix's "all three must pass" rule still applies via the -per-cell aggregator. +The three Claude tiers are probed concurrently via `run_probe_cell`, +the same way the CLI rows fan out; each probe first waits on the +shared per-provider token bucket, so a sequential cell paid that wait +three times before reporting. The matrix's "all three must pass" rule +still applies via the per-cell aggregator. """ from __future__ import annotations @@ -46,6 +46,7 @@ from __future__ import annotations import pytest from claude_code._env import require_proxy_client +from claude_code._probe_cell import run_probe_cell from claude_code.http_probe import ( assert_tool_search_shape, probe_tool_search, @@ -67,17 +68,12 @@ def test_tool_search_azure(compat_result): tier.""" client, api_key = require_proxy_client(compat_result) - failures = [] - for model in AZURE_MODELS: - result = probe_tool_search(client=client, api_key=api_key, model=model) - shape_error = assert_tool_search_shape(result) - if shape_error is not None: - error = f"[{model}] tool_search probe failed: {shape_error}" - compat_result.add({"status": "fail", "error": error}) - failures.append(error) - continue - - compat_result.add({"status": "pass"}) - - if failures: - pytest.fail("; ".join(failures), pytrace=False) + run_probe_cell( + compat_result=compat_result, + models=AZURE_MODELS, + probe=lambda model: probe_tool_search( + client=client, api_key=api_key, model=model + ), + check_shape=assert_tool_search_shape, + probe_name="tool_search", + ) diff --git a/tests/e2e/claude_code/tool_search/test_bedrock_converse.py b/tests/e2e/claude_code/tool_search/test_bedrock_converse.py index f395122a5ab..ba29c0afe02 100644 --- a/tests/e2e/claude_code/tool_search/test_bedrock_converse.py +++ b/tests/e2e/claude_code/tool_search/test_bedrock_converse.py @@ -34,11 +34,11 @@ and forwarding, and the upstream either accepts or 400s. A red cell here is always a proxy-side regression, not a flaky model-behavior artifact. -Three Claude tiers are probed in sequence (count is too low to be -worth the parallelism overhead, and HTTP probes don't compete for -the proxy's `--num-workers` slots the way CLI subprocess runs do). -The matrix's "all three must pass" rule still applies via the -per-cell aggregator. +The three Claude tiers are probed concurrently via `run_probe_cell`, +the same way the CLI rows fan out; each probe first waits on the +shared per-provider token bucket, so a sequential cell paid that wait +three times before reporting. The matrix's "all three must pass" rule +still applies via the per-cell aggregator. """ from __future__ import annotations @@ -46,6 +46,7 @@ from __future__ import annotations import pytest from claude_code._env import require_proxy_client +from claude_code._probe_cell import run_probe_cell from claude_code.http_probe import ( assert_tool_search_shape, probe_tool_search, @@ -66,17 +67,12 @@ def test_tool_search_bedrock_converse(compat_result): tier.""" client, api_key = require_proxy_client(compat_result) - failures = [] - for model in BEDROCK_CONVERSE_MODELS: - result = probe_tool_search(client=client, api_key=api_key, model=model) - shape_error = assert_tool_search_shape(result) - if shape_error is not None: - error = f"[{model}] tool_search probe failed: {shape_error}" - compat_result.add({"status": "fail", "error": error}) - failures.append(error) - continue - - compat_result.add({"status": "pass"}) - - if failures: - pytest.fail("; ".join(failures), pytrace=False) + run_probe_cell( + compat_result=compat_result, + models=BEDROCK_CONVERSE_MODELS, + probe=lambda model: probe_tool_search( + client=client, api_key=api_key, model=model + ), + check_shape=assert_tool_search_shape, + probe_name="tool_search", + ) diff --git a/tests/e2e/claude_code/tool_search/test_bedrock_invoke.py b/tests/e2e/claude_code/tool_search/test_bedrock_invoke.py index 12f8909e3e8..3b3d075a2b7 100644 --- a/tests/e2e/claude_code/tool_search/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/tool_search/test_bedrock_invoke.py @@ -34,11 +34,11 @@ and forwarding, and the upstream either accepts or 400s. A red cell here is always a proxy-side regression, not a flaky model-behavior artifact. -Three Claude tiers are probed in sequence (count is too low to be -worth the parallelism overhead, and HTTP probes don't compete for -the proxy's `--num-workers` slots the way CLI subprocess runs do). -The matrix's "all three must pass" rule still applies via the -per-cell aggregator. +The three Claude tiers are probed concurrently via `run_probe_cell`, +the same way the CLI rows fan out; each probe first waits on the +shared per-provider token bucket, so a sequential cell paid that wait +three times before reporting. The matrix's "all three must pass" rule +still applies via the per-cell aggregator. """ from __future__ import annotations @@ -46,6 +46,7 @@ from __future__ import annotations import pytest from claude_code._env import require_proxy_client +from claude_code._probe_cell import run_probe_cell from claude_code.http_probe import ( assert_tool_search_shape, probe_tool_search, @@ -70,17 +71,12 @@ def test_tool_search_bedrock_invoke(compat_result): tier.""" client, api_key = require_proxy_client(compat_result) - failures = [] - for model in BEDROCK_INVOKE_MODELS: - result = probe_tool_search(client=client, api_key=api_key, model=model) - shape_error = assert_tool_search_shape(result) - if shape_error is not None: - error = f"[{model}] tool_search probe failed: {shape_error}" - compat_result.add({"status": "fail", "error": error}) - failures.append(error) - continue - - compat_result.add({"status": "pass"}) - - if failures: - pytest.fail("; ".join(failures), pytrace=False) + run_probe_cell( + compat_result=compat_result, + models=BEDROCK_INVOKE_MODELS, + probe=lambda model: probe_tool_search( + client=client, api_key=api_key, model=model + ), + check_shape=assert_tool_search_shape, + probe_name="tool_search", + ) diff --git a/tests/e2e/claude_code/tool_search/test_vertex_ai.py b/tests/e2e/claude_code/tool_search/test_vertex_ai.py index 7d0d35b1c1d..e39a58e75ff 100644 --- a/tests/e2e/claude_code/tool_search/test_vertex_ai.py +++ b/tests/e2e/claude_code/tool_search/test_vertex_ai.py @@ -34,11 +34,11 @@ and forwarding, and the upstream either accepts or 400s. A red cell here is always a proxy-side regression, not a flaky model-behavior artifact. -Three Claude tiers are probed in sequence (count is too low to be -worth the parallelism overhead, and HTTP probes don't compete for -the proxy's `--num-workers` slots the way CLI subprocess runs do). -The matrix's "all three must pass" rule still applies via the -per-cell aggregator. +The three Claude tiers are probed concurrently via `run_probe_cell`, +the same way the CLI rows fan out; each probe first waits on the +shared per-provider token bucket, so a sequential cell paid that wait +three times before reporting. The matrix's "all three must pass" rule +still applies via the per-cell aggregator. """ from __future__ import annotations @@ -46,6 +46,7 @@ from __future__ import annotations import pytest from claude_code._env import require_proxy_client +from claude_code._probe_cell import run_probe_cell from claude_code.http_probe import ( assert_tool_search_shape, probe_tool_search, @@ -67,17 +68,12 @@ def test_tool_search_vertex_ai(compat_result): tier.""" client, api_key = require_proxy_client(compat_result) - failures = [] - for model in VERTEX_AI_MODELS: - result = probe_tool_search(client=client, api_key=api_key, model=model) - shape_error = assert_tool_search_shape(result) - if shape_error is not None: - error = f"[{model}] tool_search probe failed: {shape_error}" - compat_result.add({"status": "fail", "error": error}) - failures.append(error) - continue - - compat_result.add({"status": "pass"}) - - if failures: - pytest.fail("; ".join(failures), pytrace=False) + run_probe_cell( + compat_result=compat_result, + models=VERTEX_AI_MODELS, + probe=lambda model: probe_tool_search( + client=client, api_key=api_key, model=model + ), + check_shape=assert_tool_search_shape, + probe_name="tool_search", + )