diff --git a/tests/integration/authorization/test_audit_any_sweep_auth.py b/tests/integration/authorization/test_audit_any_sweep_auth.py index f6f5caa5240..c0142cd6f75 100644 --- a/tests/integration/authorization/test_audit_any_sweep_auth.py +++ b/tests/integration/authorization/test_audit_any_sweep_auth.py @@ -9,7 +9,7 @@ import httpx import openai import pytest import websockets -from integration._support.client import Gateway, eventually +from integration._support.client import Gateway, Scenario, eventually from integration._support.database import read_rows from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue @@ -63,7 +63,8 @@ _RESPONSES_EVENTS: Final = ( ) -def _anthropic_model(gateway: Gateway, *, stream: bool) -> str: +def _anthropic_model(scenario: Scenario, *, stream: bool) -> str: + gateway: Final = scenario.gateway scenario_id: Final = "audit-anthropic-stream" if stream else "audit-anthropic" if stream: _register_scenario(gateway, scenario_id, {"content_type": "text/event-stream", "frames": _ANTHROPIC_EVENTS}) @@ -76,23 +77,13 @@ def _anthropic_model(gateway: Gateway, *, stream: bool) -> str: "routes": {"POST /v1/messages": {"content_type": "application/json", "body": _ANTHROPIC_MESSAGE}}, }, ) - name: Final = f"integration-{uuid.uuid4().hex}" - created: Final = gateway.post( - "/model/new", - { - "model_name": name, - "litellm_params": { - "model": "anthropic/claude-3-5-sonnet-20241022", - "api_key": "integration-provider-key", - "api_base": f"{gateway.upstream_url}/{scenario_id}", - }, - }, + return scenario.model( + model="anthropic/claude-3-5-sonnet-20241022", api_base=f"{gateway.upstream_url}/{scenario_id}" ) - assert isinstance(created["model_info"], dict) and created["model_info"]["id"], created - return name -def _responses_model(gateway: Gateway, *, stream: bool) -> str: +def _responses_model(scenario: Scenario, *, stream: bool) -> str: + gateway: Final = scenario.gateway scenario_id: Final = "audit-resp-stream" if stream else "audit-resp" if stream: _register_scenario(gateway, scenario_id, {"content_type": "text/event-stream", "frames": _RESPONSES_EVENTS}) @@ -105,20 +96,7 @@ def _responses_model(gateway: Gateway, *, stream: bool) -> str: "routes": {"POST /responses": {"content_type": "application/json", "body": _RESPONSES_BODY}}, }, ) - name: Final = f"integration-{uuid.uuid4().hex}" - created: Final = gateway.post( - "/model/new", - { - "model_name": name, - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "integration-provider-key", - "api_base": f"{gateway.upstream_url}/{scenario_id}", - }, - }, - ) - assert isinstance(created["model_info"], dict) and created["model_info"]["id"], created - return name + return scenario.model(api_base=f"{gateway.upstream_url}/{scenario_id}") def _openai_client(gateway: Gateway, key: str) -> openai.OpenAI: @@ -219,8 +197,8 @@ def test_allowed_key_chat_completions_restamps_alias( @pytest.mark.parametrize("client_kind", ("httpx", "anthropic", "anthropic_async")) @pytest.mark.parametrize("stream", (False, True)) def test_allowed_key_messages(gateway: Gateway, client_kind: str, stream: bool) -> None: - model: Final = _anthropic_model(gateway, stream=stream) with gateway.scenario() as scenario: + model: Final = _anthropic_model(scenario, stream=stream) key: Final = scenario.key(models=[model]) body: Final = { "model": model, @@ -264,8 +242,8 @@ def test_allowed_key_messages(gateway: Gateway, client_kind: str, stream: bool) @pytest.mark.parametrize("client_kind", ("httpx", "openai", "openai_async")) @pytest.mark.parametrize("stream", (False, True)) def test_allowed_key_responses(gateway: Gateway, client_kind: str, stream: bool) -> None: - model: Final = _responses_model(gateway, stream=stream) with gateway.scenario() as scenario: + model: Final = _responses_model(scenario, stream=stream) key: Final = scenario.key(models=[model]) body: Final = { "model": model, @@ -334,20 +312,10 @@ async def test_realtime_websocket_allowed_and_denied(gateway: Gateway) -> None: ), }, ) - name: Final = f"integration-{uuid.uuid4().hex}" - created: Final = gateway.post( - "/model/new", - { - "model_name": name, - "litellm_params": { - "model": "openai/gpt-4o-realtime-preview", - "api_key": rt_scenario, - "api_base": gateway.upstream_url, - }, - }, - ) - assert isinstance(created["model_info"], dict) and created["model_info"]["id"], created with gateway.scenario() as scenario: + name: Final = scenario.model( + model="openai/gpt-4o-realtime-preview", api_key=rt_scenario, api_base=gateway.upstream_url + ) key: Final = scenario.key(models=[name]) other_model: Final = scenario.model() denied_key: Final = scenario.key(models=[other_model]) diff --git a/tests/integration/caching/response_cache_case.py b/tests/integration/caching/response_cache_case.py index a1850ebdb9d..312f29faeb3 100644 --- a/tests/integration/caching/response_cache_case.py +++ b/tests/integration/caching/response_cache_case.py @@ -279,16 +279,19 @@ def redis_store() -> Redis: return Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) -def entry_key(store: Redis, needle: str) -> str | None: - """The response cache key (a bare sha256) whose stored value names ``needle``, or ``None``.""" +def entry_keys(store: Redis, needle: str) -> Iterator[str]: for raw_key in store.scan_iter(count=1000): key: Final = raw_key.decode() if not CACHE_KEY.match(key) or store.type(key) != b"string": continue value: Final = store.get(key) if value is not None and needle.encode() in value: - return key - return None + yield key + + +def entry_key(store: Redis, needle: str) -> str | None: + """The response cache key (a bare sha256) whose stored value names ``needle``, or ``None``.""" + return next(entry_keys(store, needle), None) def await_entry(store: Redis, needle: str) -> str: diff --git a/tests/integration/caching/test_cache_max_messages.py b/tests/integration/caching/test_cache_max_messages.py index 32c71e887fe..e28015a5ba4 100644 --- a/tests/integration/caching/test_cache_max_messages.py +++ b/tests/integration/caching/test_cache_max_messages.py @@ -15,6 +15,7 @@ from integration._support.openai_wire import chat_reply, responses_reply from integration._support.process import owned_proxy from integration._support.provider import SharedProvider from integration._support.wire import Reply, Request, wire_server +from integration.caching.response_cache_case import await_entry, redis_store from pydantic import JsonValue, TypeAdapter from redis import Redis from redis.exceptions import ResponseError @@ -133,11 +134,6 @@ def _first(value: JsonValue) -> dict[str, JsonValue]: return object_value(value[0]) -def _cached_responses(redis: Redis[bytes]) -> frozenset[bytes]: - digests: Final = tuple(key for key in redis.scan_iter() if len(key) == 64) - return frozenset(key for key in digests if b'"response"' in (redis.get(key) or b"")) - - _SCENARIOS: Final = ( pytest.param("/v1/messages", "anthropic/claude-sonnet-5-5", _claude_code_turns, id="messages"), pytest.param("/v1/chat/completions", "openai/gpt-5.4", _agent_turns, id="chat-completions"), @@ -149,22 +145,22 @@ _SCENARIOS: Final = ( def test_cache_serves_a_turn_under_max_messages_and_skips_one_past_it( gateway: Gateway, provider: SharedProvider, path: str, model: str, turns: _Turns ) -> None: - under_cap, past_cap = turns(f"update the config {uuid.uuid4().hex}") - provider.expect(_reply(path, "first answer"), _reply(path, "second answer"), _reply(path, "third answer")) + task: Final = uuid.uuid4().hex + under_cap, past_cap = turns(f"update the config {task}") + answers: Final = tuple(f"{ordinal} answer {task}" for ordinal in ("first", "second", "third")) + provider.expect(*(_reply(path, answer) for answer in answers)) - redis_client: Final[Redis[bytes]] = Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) - with redis_client as redis: - cached_before: Final = _cached_responses(redis) + with redis_store() as store: gateway.post(path, _body(path, model, under_cap)) - eventually(lambda: _cached_responses(redis) - cached_before, lambda written: len(written) == 1) + await_entry(store, answers[0]) repeated: Final = gateway.post(path, _body(path, model, under_cap)) past_cap_twice: Final = ( gateway.post(path, _body(path, model, past_cap)), gateway.post(path, _body(path, model, past_cap)), ) - assert _answer(path, repeated) == "first answer", "a repeated turn under max_messages was not served from the cache" - assert tuple(_answer(path, answer) for answer in past_cap_twice) == ("second answer", "third answer"), ( + assert _answer(path, repeated) == answers[0], "a repeated turn under max_messages was not served from the cache" + assert tuple(_answer(path, answer) for answer in past_cap_twice) == answers[1:], ( "a turn past max_messages was served from the cache" ) assert len(provider.received()) == 3 diff --git a/tests/integration/observability/test_s3_v2_flush_surfaces.py b/tests/integration/observability/test_s3_v2_flush_surfaces.py index 1e373171de6..5a94eaa248f 100644 --- a/tests/integration/observability/test_s3_v2_flush_surfaces.py +++ b/tests/integration/observability/test_s3_v2_flush_surfaces.py @@ -1,11 +1,9 @@ -import os import re import uuid from pathlib import Path from typing import Final import pytest -from redis import Redis from _s3_v2_support import ( BUCKET, PREFIX, @@ -21,6 +19,7 @@ from _s3_v2_support import ( from integration._support.client import Gateway, eventually from integration._support.process import owned_proxy from integration._support.wire import wire_server +from integration.caching.response_cache_case import entry_keys, redis_store pytestmark: Final = pytest.mark.shares_worker_state @@ -117,13 +116,14 @@ def test_s3_v2_cache_hit_twins_log_one_object_per_request(gateway: Gateway, tmp_ model="anthropic/claude-sonnet-4-5-20250929", api_base=provider.url, api_key="synthetic-provider-key" ) key: Final = scenario.key(models=[openai_model, anthropic_model]) - cache: Final = Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) - keys_before: Final = cache.dbsize() warmed: Final = tuple( call_surface(candidate, surface, openai_model, anthropic_model, key, f"{marker}-{surface}") for surface in SURFACES ) - eventually(cache.dbsize, lambda size: size >= keys_before + len(SURFACES), seconds=30) + with redis_store() as cache: + eventually( + lambda: frozenset(entry_keys(cache, marker)), lambda owned: len(owned) >= len(SURFACES), seconds=30 + ) repeated: Final = tuple( call_surface(candidate, surface, openai_model, anthropic_model, key, f"{marker}-{surface}", False) for surface in SURFACES diff --git a/tests/integration/translation/chat_completions/reasoning/test_chat_completions_reasoning_gemini.py b/tests/integration/translation/chat_completions/reasoning/test_chat_completions_reasoning_gemini.py index efd649f9383..e2369fea443 100644 --- a/tests/integration/translation/chat_completions/reasoning/test_chat_completions_reasoning_gemini.py +++ b/tests/integration/translation/chat_completions/reasoning/test_chat_completions_reasoning_gemini.py @@ -8,11 +8,11 @@ from unittest.mock import ANY from uuid import uuid4 import pytest -from integration._support.client import Gateway, eventually, list_value, object_value, string_value +from integration._support.client import Gateway, list_value, object_value, string_value from integration._support.process import graceful_stop_seconds from integration._support.provider import SharedProvider from integration._support.wire import Reply -from integration.caching.response_cache_case import CACHE_KEY, redis_store +from integration.caching.response_cache_case import await_entry, redis_store from integration.translation.case import SSEFrame, TranslationTestCase from integration.translation.chat_completions.bases.gemini import GEMINI_3_8_FLASH_TEST_CASE from integration.translation.chat_completions.bases.vertex_ai import ( @@ -840,7 +840,7 @@ def without_closed_schema(tool: JsonValue) -> JsonValue: return {**object_value(tool), "function": {**function, "parameters": open_parameters}} -def cacheable_run(case: TranslationTestCase, marker: str) -> TranslationTestCase: +def cacheable_run(case: TranslationTestCase, marker: str, reply_id: str) -> TranslationTestCase: messages: Final = list_value(case.litellm_request["messages"]) system: Final = object_value(messages[0]) instruction: Final = object_value(case.expected_provider_request["system_instruction"]) @@ -853,6 +853,7 @@ def cacheable_run(case: TranslationTestCase, marker: str) -> TranslationTestCase "messages": [{**system, "content": text}, messages[1]], "tools": [without_closed_schema(tool) for tool in list_value(case.litellm_request["tools"])], }, + mock_provider_response={**case.mock_provider_response, "responseId": reply_id}, expected_provider_request={ **case.expected_provider_request, "contents": [list_value(case.expected_provider_request["contents"])[0]], @@ -861,14 +862,10 @@ def cacheable_run(case: TranslationTestCase, marker: str) -> TranslationTestCase ) -def response_cache_keys() -> frozenset[str]: - keys: Final = (raw.decode() for raw in redis_store().scan_iter()) - return frozenset(key for key in keys if CACHE_KEY.match(key)) - - -def stored_reply_key(before: frozenset[str]) -> str: - (key,) = eventually(lambda: response_cache_keys() - before, lambda added: len(added) == 1) - return key +def stored_reply_key(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> str: + reply: Final = assert_translation(case, gateway, provider) + with redis_store() as store: + return await_entry(store, string_value(reply["id"])) GEMINI_3_8_FLASH_SUMMARY_ALIAS_TEST_CASE: Final = with_request( @@ -888,12 +885,11 @@ def test_a_summary_alias_ask_and_a_plain_request_never_share_a_cached_reply( first: TranslationTestCase, second: TranslationTestCase, thoughts_off: Gateway, provider: SharedProvider ) -> None: marker: Final = uuid4().hex - before: Final = response_cache_keys() - assert_translation(cacheable_run(first, marker), thoughts_off, provider) - first_key: Final = stored_reply_key(before) - assert_translation(cacheable_run(second, marker), thoughts_off, provider) - second_key: Final = stored_reply_key(before | {first_key}) - repeat: Final = thoughts_off.request("POST", second.litellm_endpoint, cacheable_run(second, marker).litellm_request) + second_run: Final = cacheable_run(second, marker, uuid4().hex) + first_key: Final = stored_reply_key(cacheable_run(first, marker, uuid4().hex), thoughts_off, provider) + second_key: Final = stored_reply_key(second_run, thoughts_off, provider) + assert second_key != first_key + repeat: Final = thoughts_off.request("POST", second.litellm_endpoint, second_run.litellm_request) assert repeat.status_code == 200, repeat.text assert repeat.headers.get("x-litellm-cache-key") == second_key, sorted(repeat.headers) assert provider.received() == ()