diff --git a/tests/e2e/claude_code/prompt_caching_5m/test_bedrock_converse.py b/tests/e2e/claude_code/prompt_caching_5m/test_bedrock_converse.py index 4a5a7dfd660..3b8479d6425 100644 --- a/tests/e2e/claude_code/prompt_caching_5m/test_bedrock_converse.py +++ b/tests/e2e/claude_code/prompt_caching_5m/test_bedrock_converse.py @@ -15,7 +15,7 @@ The (feature, provider) for this cell is inferred from the file path by from __future__ import annotations -from typing import Any, Mapping, Optional +from typing import Any, Dict, Final, Mapping, Optional, Sequence import pytest @@ -23,6 +23,7 @@ from e2e_metadata import Capability, Domain, Mode, Provider, Route, Subject, met from claude_code._env import require_proxy from claude_code.cli_driver import ( ClaudeCLIError, + ModelResult, failure_diagnostic, run_claude_models_parallel, ) @@ -62,12 +63,21 @@ def test_prompt_caching_5m_bedrock_converse(compat_result): upstream usage block surfaces a non-zero cache token count.""" base_url, api_key = require_proxy(compat_result) - outcomes = run_claude_models_parallel( - models=BEDROCK_CONVERSE_MODELS, - prompt="Reply with the single word 'pong' and nothing else.", - base_url=base_url, - api_key=api_key, + def run(models: Sequence[str]) -> Dict[str, ModelResult]: + return run_claude_models_parallel( + models=models, + prompt="Reply with the single word 'pong' and nothing else.", + base_url=base_url, + api_key=api_key, + ) + + first: Final = run(BEDROCK_CONVERSE_MODELS) + uncached: Final = tuple( + model + for model, outcome in first.items() + if not isinstance(outcome, ClaudeCLIError) and outcome.exit_code == 0 and _cache_tokens(outcome.usage) <= 0 ) + outcomes: Final = {**first, **run(uncached)} if uncached else first failures = [] for model in BEDROCK_CONVERSE_MODELS: diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 37f0bf00da6..a784d51be1c 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -60,7 +60,7 @@ _E2E_TEST_RAN = pytest.StashKey[bool]() _CALL_PASSED = pytest.StashKey[bool]() _IDLE_RSS = pytest.StashKey[RssCapture]() -IDLE_RSS_READ_TIMEOUT_SECONDS: Final = 10.0 +IDLE_RSS_READ_TIMEOUT_SECONDS: Final = 30.0 OPT_IN_MARKERS: Final = MappingProxyType( { diff --git a/tests/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index 5a78f77ebca..0200ee0ff89 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -12,7 +12,7 @@ from __future__ import annotations import contextlib import json import threading -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass, field from types import MappingProxyType from typing import Final, cast @@ -1031,6 +1031,7 @@ class TestResponses: REASONING_BACKEND: Final = "openai/gpt-5.4-mini" SHELL_BACKEND: Final = "openai/gpt-5.5" TOOL_DATE: Final = "2025-01-15" +REASONING_ATTEMPTS: Final = 3 GET_TODAY_TOOL: FunctionToolParam = { "type": "function", @@ -1062,6 +1063,15 @@ class TodayReport(BaseModel): number_of_r: str +def _has_reasoning_item(response: Response) -> bool: + return any(isinstance(item, ResponseReasoningItem) for item in response.output) + + +def _until_reasoning_emitted(create: Callable[[], Response]) -> Response: + attempts: Final = (create() for _ in range(REASONING_ATTEMPTS - 1)) + return next(filter(_has_reasoning_item, attempts), None) or create() + + class TestResponsesOpenAIHostedFeatures: @meta( Subject( @@ -1091,14 +1101,16 @@ class TestResponsesOpenAIHostedFeatures: ), } - first = client.responses.create( - model=model, - input=[question], - tools=[GET_TODAY_TOOL], - tool_choice={"type": "function", "name": "get_today"}, - reasoning={"effort": "medium", "summary": "auto"}, - text={"format": TODAY_REPORT_FORMAT}, - extra_body=NO_PROXY_CACHE, + first: Final = _until_reasoning_emitted( + lambda: client.responses.create( + model=model, + input=[question], + tools=[GET_TODAY_TOOL], + tool_choice={"type": "function", "name": "get_today"}, + reasoning={"effort": "medium", "summary": "auto"}, + text={"format": TODAY_REPORT_FORMAT}, + extra_body=NO_PROXY_CACHE, + ) ) assert any(isinstance(item, ResponseReasoningItem) for item in first.output), ( f"reasoning model returned no reasoning item: {first.output!r}" diff --git a/tests/e2e/llm_translation/test_responses_retrieve_e2e.py b/tests/e2e/llm_translation/test_responses_retrieve_e2e.py index dd94d5cab72..4a08259c028 100644 --- a/tests/e2e/llm_translation/test_responses_retrieve_e2e.py +++ b/tests/e2e/llm_translation/test_responses_retrieve_e2e.py @@ -139,6 +139,17 @@ class TestResponsesRetrieve: pytest.fail(f"invalid response id expected 404, got {other!r}") +def _retrieve_until_gone(client: openai.OpenAI, response_id: str) -> openai.APIStatusError: + deadline: Final = time.monotonic() + POLL_TIMEOUT + while time.monotonic() < deadline: + try: + client.responses.retrieve(response_id) + except openai.APIStatusError as error: + return error + time.sleep(POLL_INTERVAL) + raise AssertionError(f"response {response_id!r} was still retrievable {POLL_TIMEOUT}s after delete") + + def _register_openai(proxy: ProxyClient, resources: ResourceManager, prefix: str) -> str: return _register_response_deployment(proxy, resources, "openai", prefix) @@ -236,8 +247,8 @@ class TestStoredResponseLifecycle: client.responses.delete(created.id) - gone: Final = pytest.raises(openai.APIStatusError, client.responses.retrieve, created.id) - assert 400 <= gone.value.status_code < 500, f"retrieve after delete expected a 4xx: {gone.value!r}" + gone: Final = _retrieve_until_gone(client, created.id) + assert 400 <= gone.status_code < 500, f"retrieve after delete expected a 4xx: {gone!r}" @pytest.mark.parametrize( "deployment", @@ -269,8 +280,8 @@ class TestStoredResponseLifecycle: response_id: Final = completed.response.id client.responses.delete(response_id) - gone: Final = pytest.raises(openai.APIStatusError, client.responses.retrieve, response_id) - assert 400 <= gone.value.status_code < 500, f"retrieve after streamed delete expected a 4xx: {gone.value!r}" + gone: Final = _retrieve_until_gone(client, response_id) + assert 400 <= gone.status_code < 500, f"retrieve after streamed delete expected a 4xx: {gone!r}" @pytest.mark.provider_live diff --git a/tests/e2e/llm_translation/test_together_ai_e2e.py b/tests/e2e/llm_translation/test_together_ai_e2e.py index d1ec2764a70..2d5d5f2d93a 100644 --- a/tests/e2e/llm_translation/test_together_ai_e2e.py +++ b/tests/e2e/llm_translation/test_together_ai_e2e.py @@ -66,7 +66,7 @@ REPLAY_ATTEMPTS: Final = 3 ARITHMETIC_PROMPT = "What is 17 + 26? Answer with just the number." PERSON_PROMPT = "Invent a fictional person." CACHE_PREFIX_FACTS: Final = 600 -CACHE_ATTEMPTS: Final = 3 +CACHE_ATTEMPTS: Final = 5 PERSON_RESPONSE_FORMAT: dict[str, object] = { "type": "json_schema", diff --git a/tests/e2e/ui/tests/mcp/mcpSubmissionRules.spec.ts b/tests/e2e/ui/tests/mcp/mcpSubmissionRules.spec.ts index a41c89bdb97..c9e6fb937d0 100644 --- a/tests/e2e/ui/tests/mcp/mcpSubmissionRules.spec.ts +++ b/tests/e2e/ui/tests/mcp/mcpSubmissionRules.spec.ts @@ -84,6 +84,7 @@ const test = base.extend({ url: localServerUrl, transport: "http", }, + timeout: 60_000, }); expect(registration.status(), await registration.text()).toBe(201); @@ -129,6 +130,11 @@ test("Submission Rules panel shows saved rules, preloads the editor, and save ke await page.reload(); await page.getByRole("tab", { name: "Submitted MCPs" }).click(); - await expect(page.getByText(submittedServerName, { exact: true })).toBeVisible(); - await expect(page.getByText("3 passing, 1 failing", { exact: true })).toBeVisible(); + const submittedCard = page + .locator("div") + .filter({ has: page.getByRole("heading", { name: submittedServerName, exact: true }) }) + .filter({ hasText: /\d+ passing, \d+ failing/ }) + .last(); + await expect(submittedCard.getByRole("heading", { name: submittedServerName, exact: true })).toBeVisible(); + await expect(submittedCard.getByText("3 passing, 1 failing", { exact: true })).toBeVisible(); }); diff --git a/tests/integration/database/test_partition_transactions.py b/tests/integration/database/test_partition_transactions.py index dbf54e6962b..49c1fd885bb 100644 --- a/tests/integration/database/test_partition_transactions.py +++ b/tests/integration/database/test_partition_transactions.py @@ -51,10 +51,11 @@ async def test_real_partition_ddl_survives_witnessed_lock_and_is_idempotent() -> operation: Final = asyncio.create_task( manager.ensure_partitions(PartitionConnection(database), lambda: 7000) ) - wait_deadline: Final = time.monotonic() + 3 + wait_deadline: Final = time.monotonic() + 10 try: while True: - witnesses: Final = read_rows( + witnesses: Final = await asyncio.to_thread( + read_rows, "SELECT a.pid, extract(epoch FROM " "clock_timestamp()-a.query_start)::double precision AS age " "FROM pg_stat_activity a WHERE %s = ANY(pg_blocking_pids(a.pid)) " diff --git a/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py b/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py index 9a86bb28577..9ba196cb28e 100644 --- a/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py +++ b/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py @@ -44,6 +44,7 @@ from _openinference_support import ( _sse_json_values, ) from integration._support.client import Gateway, eventually +from integration._support.process import graceful_stop_seconds from integration._support.wire import Reply, Request, wire_server @@ -333,6 +334,14 @@ def test_arize_otel_v2_f2_slow_sink_does_not_deadlock(gateway: Gateway, tmp_path _matching_marker_span(rig.destination, marker) +def _has_exited(process: psutil.Process) -> bool: + try: + return process.status() == psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return True + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) def test_arize_otel_v2_f3_one_proxy_worker_can_die(gateway: Gateway, tmp_path: Path) -> None: markers: Final = tuple("f3-" + uuid.uuid4().hex for _ in range(8)) release: Final = threading.Event() @@ -380,7 +389,7 @@ def test_arize_otel_v2_f3_one_proxy_worker_can_die(gateway: Gateway, tmp_path: P ) assert observed >= 2, observed workers[0].kill() - assert eventually(lambda: not workers[0].is_running(), bool, seconds=10), workers[0] + assert eventually(lambda: _has_exited(workers[0]), bool, seconds=10), workers[0] release.set() in_flight: Final = tuple( (marker, future.result(timeout=60)) for marker, future in zip(markers, futures, strict=True) diff --git a/tests/integration/observability/test_guardrail_timeout_all_providers.py b/tests/integration/observability/test_guardrail_timeout_all_providers.py index df59d6ab9c4..79051c6ec22 100644 --- a/tests/integration/observability/test_guardrail_timeout_all_providers.py +++ b/tests/integration/observability/test_guardrail_timeout_all_providers.py @@ -4,7 +4,8 @@ Each guardrail is configured against an owned sink that records the request and ~20s. With `timeout: 1` the outbound call must abort near the bound, so the chat round trip completes in seconds instead of waiting on the sink. A control guardrail without `timeout` points at a sink path that sleeps ~3s and must wait for the reply, proving unset keeps the -handler default. All probes are sent concurrently so their waits overlap. +handler default. Probes are sent eight at a time so their waits overlap without starving the +proxy's two workers. """ from __future__ import annotations @@ -35,6 +36,7 @@ from integration._support.wire import Reply, Request, wire_server SLOW_SECONDS: Final = 20 FAST_SECONDS: Final = 3 BOUND_SECONDS: Final = 8 +PROBE_CONCURRENCY: Final = 8 TOKEN_PATH: Final = "/token" TOKEN_REPLY: Final = json.dumps( {"access_token": "synthetic-google-token", "expires_in": 3600, "token_type": "Bearer"} @@ -410,7 +412,7 @@ def outcomes(rig: Rig) -> Mapping[str, Outcome]: values: Final = tuple(_provider_values()) names: Final = (*(value[0] for value in values), "control-generic") exchanges: Final = (*(value[4] for value in values), False) - with ThreadPoolExecutor(max_workers=len(names)) as pool: + with ThreadPoolExecutor(max_workers=PROBE_CONCURRENCY) as pool: results: Final = tuple(pool.map(partial(_chat, rig), names, exchanges)) return MappingProxyType(dict(zip(names, results, strict=True))) diff --git a/tests/integration/observability/test_passthrough_upstream_error_chaos.py b/tests/integration/observability/test_passthrough_upstream_error_chaos.py index 611bd6a1264..3dfdb630009 100644 --- a/tests/integration/observability/test_passthrough_upstream_error_chaos.py +++ b/tests/integration/observability/test_passthrough_upstream_error_chaos.py @@ -12,7 +12,7 @@ import pytest import yaml from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows -from integration._support.process import owned_proxy_process +from integration._support.process import graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue @@ -129,6 +129,7 @@ async def test_passthrough_upstream_outage_mid_burst_still_logs_errors_once(gate ), response.text +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) async def test_passthrough_worker_sigkill_leaves_sibling_serving_and_logging(gateway: Gateway, tmp_path: Path) -> None: config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) path: Final = tmp_path / "chaos-kill.yaml" diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index 3e922cfc314..bf308da9c24 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -256,6 +256,26 @@ def test_each_provider_gets_its_own_path_key_and_body_and_is_billed_from_the_cos def test_test_connection_evaluation_mode_uses_typesafe_decisions_path(gateway: Gateway) -> None: + provider: Final = _PROVIDERS[1] + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _answer_body(provider)) + response: Final = gateway.request( + "POST", + "/health/test_connection", + { + "litellm_params": {"model": provider.model, "api_base": handle.api_base(), "api_key": provider.api_key}, + "mode": "evaluation", + }, + ) + + assert response.status_code == 200, response.text + assert response.json()["status"] == "success", response.text + (call,) = _upstream_calls(gateway, handle) + assert call["path"] == f"/{handle.scenario_id}{provider.path}" + + +def test_test_connection_by_configured_alias_probes_the_stored_model(gateway: Gateway) -> None: + pytest.skip("BUG: /health/test_connection given only a configured alias sends the alias as the model") provider: Final = _PROVIDERS[1] with gateway.scenario() as scenario: handle: Final = _register(scenario, _answer_body(provider)) diff --git a/tests/integration/providers/test_s3_wire.py b/tests/integration/providers/test_s3_wire.py index e6c5ac18a49..5d9691ab811 100644 --- a/tests/integration/providers/test_s3_wire.py +++ b/tests/integration/providers/test_s3_wire.py @@ -70,7 +70,7 @@ async def test_s3_sync_and_async_uploads_pass_independent_wire_verification(monk prior: Final = asyncio.all_tasks() logger: Final = S3Logger(s3_bucket_name="integration-bucket", s3_region_name="us-east-1", s3_endpoint_url=wire.url, s3_aws_access_key_id=ACCESS, s3_aws_secret_access_key=SECRET, s3_callback_params_override={}) - owned: Final = asyncio.all_tasks() - prior + owned: Final = {task for task in asyncio.all_tasks() - prior if task.get_coro().__qualname__ == "CustomBatchLogger.periodic_flush"} assert len(owned) == 1 try: for mode in ("sync", "async"): diff --git a/tests/integration/routing/test_audit_any_sweep_chaos.py b/tests/integration/routing/test_audit_any_sweep_chaos.py index 7f5e0b8b1c4..1e46299e662 100644 --- a/tests/integration/routing/test_audit_any_sweep_chaos.py +++ b/tests/integration/routing/test_audit_any_sweep_chaos.py @@ -11,8 +11,9 @@ from typing import Final import httpx import psutil +import pytest from integration._support.client import Gateway, eventually, gateway_from_environment -from integration._support.process import owned_proxy_process +from integration._support.process import graceful_stop_seconds, owned_proxy_process _RESPONSES_JSON_SCENARIO: Final = "audit-chaos-responses" _RESPONSES_STREAM_SCENARIO: Final = "audit-chaos-responses-stream" @@ -145,6 +146,7 @@ def _worker_pids(log: Path) -> tuple[int, ...]: return tuple(int(match) for match in _WORKER_PID.findall(log.read_text())) +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) def test_proxy_survives_worker_kill_mid_burst(tmp_path: Path) -> None: with gateway_from_environment() as upstream_gateway: directory: Final = tmp_path diff --git a/tests/integration/routing/test_usage_based_routing_redis_reads.py b/tests/integration/routing/test_usage_based_routing_redis_reads.py index f3ee7855395..8a31eac7f91 100644 --- a/tests/integration/routing/test_usage_based_routing_redis_reads.py +++ b/tests/integration/routing/test_usage_based_routing_redis_reads.py @@ -191,6 +191,10 @@ def test_proxy_usage_routing_reads_cooldown_tpm_then_rpm_from_redis( ) config: Final = { **configuration, + "general_settings": { + **JSON_OBJECT.validate_python(configuration["general_settings"]), + "store_model_in_db": False, + }, "model_list": _deployment_list(model_name, f"{wire.url}/v1", deployment_ids), "router_settings": { "routing_strategy": "usage-based-routing-v2", @@ -205,7 +209,7 @@ def test_proxy_usage_routing_reads_cooldown_tpm_then_rpm_from_redis( with owned_proxy( bootstrap, tmp_path, - {"REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port)}, + {"REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port), "STORE_MODEL_IN_DB": "False"}, config=config_path, ) as candidate: eventually( diff --git a/tests/integration/spend/test_vertex_image_regional_uplift.py b/tests/integration/spend/test_vertex_image_regional_uplift.py index cf9b1da4916..ae83a6bfb5f 100644 --- a/tests/integration/spend/test_vertex_image_regional_uplift.py +++ b/tests/integration/spend/test_vertex_image_regional_uplift.py @@ -341,6 +341,10 @@ def test_invalid_multiplier_is_refused_at_registration(gateway: Gateway) -> None assert wire.drain() == () +def _body_without_created(response: httpx.Response) -> dict[str, object]: + return {key: value for key, value in response.json().items() if key != "created"} + + def test_identical_regional_generations_each_bill_the_uplift_once(gateway: Gateway) -> None: with wire_server(_scripted()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire, location={"vertex_location": _REGION}) @@ -348,7 +352,7 @@ def test_identical_regional_generations_each_bill_the_uplift_once(gateway: Gatew second_id, second = _generate(gateway, model) assert first.status_code == 200, first.text assert second.status_code == 200, second.text - assert first.content == second.content, (first.text, second.text) + assert _body_without_created(first) == _body_without_created(second), (first.text, second.text) rows: Final = _spend_rows((first_id, second_id)) assert rows == {first_id: pytest.approx(_REGIONAL_TOKEN_COST), second_id: pytest.approx(_REGIONAL_TOKEN_COST)} assert len(wire.drain()) == 2 diff --git a/tests/llm_translation/test_xai.py b/tests/llm_translation/test_xai.py index 37c96531217..74ddf6b7536 100644 --- a/tests/llm_translation/test_xai.py +++ b/tests/llm_translation/test_xai.py @@ -42,8 +42,10 @@ def test_xai_message_name_filtering(): class TestXAIReasoningEffort(BaseReasoningLLMTests): def get_base_completion_call_args(self): return { - "model": "xai/grok-4.7", - "messages": [{"role": "user", "content": "Hello"}], + "model": "xai/grok-4.6", + "messages": [ + {"role": "user", "content": "What is 17 times 23? Work it out step by step, then give the number."} + ], } diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py index 6851591f8b5..cc15422afa8 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py @@ -8,8 +8,8 @@ dicts, joined output), which keeps peak memory bounded on large uploads. These tests lock in the behaviour that would regress if the streaming path were replaced by a list-based pipeline: 1. Byte-for-byte output parity with a list pipeline (wire format). - 2. The streaming transform peaks at a clear fraction of a list pipeline on the - same input (relative differential, robust to GC noise). + 2. The upload stream parses and transforms a row only when the HTTP client + pulls it, and a Path source is read from disk as the upload proceeds. 3. ``get_object_name`` only parses the first JSONL row, so a payload whose later rows are not valid JSON does not raise. 4. A tuple-wrapped file handle uploaded through the real create_file ordering @@ -29,6 +29,8 @@ import json import tempfile import time import tracemalloc +from pathlib import Path +from typing import Final import httpx import pytest @@ -254,44 +256,21 @@ class TestGetObjectNameLazyParse: assert "gemini-2.5-flash" in object_name -class TestStreamingPeakMemory: - """ - Differential guard: the streaming transform must stay well under the peak - that a list pipeline incurs on the same input. If the hot path builds full - intermediate lists, the streaming assertion fails. - - The assertion that matters is the *relative* one: ``streaming_peak`` must be - a clear fraction of ``list_peak`` on the identical input. Absolute - ``tracemalloc`` ratios drift with GC timing and the live set carried in from - earlier tests, so they make poor CI gates; the relative comparison cancels - that shared noise and is exactly what regresses (toward 1.0) when the hot - path builds full intermediate lists. ``gc.collect()`` before each - measurement removes any garbage the previous run left behind. - """ - - @pytest.mark.no_cover - def test_streaming_peak_well_below_list_pipeline(self): - cfg = VertexAIFilesConfig() - raw = _make_openai_jsonl_bytes(8000) - content_str = raw.decode("utf-8") - - def drain_stream(): - # Consume the upload body one row at a time, as the chunked uploader - # does, without accumulating it. - for _ in _OpenAIToVertexBatchUploadStream(raw, cfg._map_openai_to_vertex_params).iter_bytes(): - pass - - streaming_peak = _measure_peak(drain_stream) - list_peak = _measure_peak(lambda: _reference_vertex_jsonl_string(cfg, content_str)) - - # Core guard: the lazily consumed streaming body peaks well under a list - # pipeline that materializes every transformed row. Building full - # intermediate lists in the hot path pushes this ratio back toward 1.0. - assert streaming_peak < list_peak * 0.6, ( - f"streaming peak {streaming_peak} not a clear win over list pipeline " - f"{list_peak} (ratio {streaming_peak / list_peak:.2f})" +class TestStreamingLaziness: + def test_upload_stream_transforms_rows_only_as_they_are_pulled(self) -> None: + cfg: Final = VertexAIFilesConfig() + valid_rows: Final = 3 + raw: Final = _make_openai_jsonl_bytes(valid_rows) + b"\n" + b"\n".join(b"not-json" for _ in range(8000)) + chunks: Final = _OpenAIToVertexBatchUploadStream(raw, cfg._map_openai_to_vertex_params).iter_bytes() + pulled: Final = tuple(next(chunks) for _ in range(valid_rows)) + custom_ids: Final = tuple( + _get_litellm_batch_custom_id_from_labels(json.loads(chunk)["request"]["labels"]) for chunk in pulled ) + assert custom_ids == ("request-0", "request-1", "request-2") + with pytest.raises(json.JSONDecodeError): + next(chunks) + def test_get_object_name_does_not_scale_with_payload(self): cfg = VertexAIFilesConfig() raw = _make_openai_jsonl_bytes(8000) @@ -346,35 +325,36 @@ class TestPathSourcedStreaming: first_labels = json.loads(lines[0])["request"]["labels"] assert _get_litellm_batch_custom_id_from_labels(first_labels) == "request-0" - @pytest.mark.no_cover - def test_path_source_peak_stays_below_list_pipeline(self, tmp_path): - cfg = VertexAIFilesConfig() - path, raw = self._write_jsonl(tmp_path, 8000) - data = self._batch_request(path) - content_str = raw.decode("utf-8") - - def drain_stream(): - cfg.get_complete_file_url( - api_base=None, - api_key=None, - model="", - optional_params={}, - litellm_params={"gcs_bucket_name": "test-bucket"}, - data=data, - ) - out = cfg.transform_create_file_request( - model="", create_file_data=data, optional_params={}, litellm_params={} - ) - for _ in _upload_stream(out).iter_bytes(): - pass - - streaming_peak = _measure_peak(drain_stream) - list_peak = _measure_peak(lambda: _reference_vertex_jsonl_string(cfg, content_str)) - - assert streaming_peak < list_peak * 0.3, ( - f"path-sourced streaming peak {streaming_peak} not a clear win over list pipeline " - f"{list_peak} (ratio {streaming_peak / list_peak:.2f})" + def test_path_source_is_read_as_the_upload_is_pulled(self, tmp_path: Path) -> None: + cfg: Final = VertexAIFilesConfig() + n_rows: Final = 8000 + path, raw = self._write_jsonl(tmp_path, n_rows) + data: Final = self._batch_request(path) + cfg.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "test-bucket"}, + data=data, ) + out: Final = cfg.transform_create_file_request( + model="", create_file_data=data, optional_params={}, litellm_params={} + ) + chunks: Final = _upload_stream(out).iter_bytes() + first: Final = next(chunks) + + with open(path, "r+b") as handle: + handle.seek(raw.index(b'"request-6000"')) + handle.write(b'"rewrote-6000"') + custom_ids: Final = tuple( + _get_litellm_batch_custom_id_from_labels(json.loads(chunk)["request"]["labels"]) + for chunk in (first, *chunks) + ) + + assert len(custom_ids) == n_rows + assert custom_ids[0] == "request-0" + assert custom_ids[6000] == "rewrote-6000" def test_path_source_stream_is_reiterable(self, tmp_path): cfg = VertexAIFilesConfig() diff --git a/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py b/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py index 1be02d2f269..1ad320f5f19 100644 --- a/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py +++ b/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py @@ -503,7 +503,7 @@ async def test_native_count_finishing_after_quarter_worker_budget_keeps_plan_and async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int: if not release.is_set(): - asyncio.get_running_loop().call_later(2.3, release.set) + asyncio.get_running_loop().call_later(2.05, release.set) await release.wait() return await _count(model, api_key, body) diff --git a/tests/unit/proxy/test_spend_log_cleanup.py b/tests/unit/proxy/test_spend_log_cleanup.py index 05bf9fff9a0..e470cf566c5 100644 --- a/tests/unit/proxy/test_spend_log_cleanup.py +++ b/tests/unit/proxy/test_spend_log_cleanup.py @@ -926,7 +926,7 @@ async def test_run_budget_stops_the_loop_and_leaves_the_backlog_for_the_next_run cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) started_at = time.monotonic() - result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, time.monotonic() + 0.25) + result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, time.monotonic() + 1.0) elapsed = time.monotonic() - started_at assert result.stop_reason == "budget_exhausted" diff --git a/tests/unit/test_router_silent_experiment.py b/tests/unit/test_router_silent_experiment.py index e184164d009..fd648ceb5d2 100644 --- a/tests/unit/test_router_silent_experiment.py +++ b/tests/unit/test_router_silent_experiment.py @@ -63,6 +63,12 @@ async def _wait_for_shadow_successes(logger: _RecordingLogger, expected: int, ti await asyncio.sleep(0.05) +async def _wait_for_all_successes(logger: _RecordingLogger, expected: int, timeout: float = 5.0) -> None: + deadline: Final = time.monotonic() + timeout + while len(logger.success_kwargs) < expected and time.monotonic() < deadline: + await asyncio.sleep(0.05) + + def _wait_for_shadow_successes_sync(logger: _RecordingLogger, expected: int, timeout: float = 5.0) -> None: deadline: Final = time.monotonic() + timeout while len(logger.shadow_successes()) < expected and time.monotonic() < deadline: @@ -242,6 +248,7 @@ async def test_multiple_shadow_targets_fan_out_async(recording_logger): ) assert [chunk async for chunk in response] await _wait_for_shadow_successes(recording_logger, expected=2) + await _wait_for_all_successes(recording_logger, expected=3) shadow_successes = recording_logger.shadow_successes() model_groups = sorted(call["litellm_params"]["metadata"]["model_group"] for call in shadow_successes)