mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
fix: cap cold MCP schema workers at usable CPUs and deflake wall-clock unit tests (#45791)
* test: inject clocks and wait on conditions in three wall-clock-flaky unit tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: keep generic API logger tests on 3.10 and cover automatic sending Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: pin router usage minute in shared rate-limit header test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: stabilize three wall-clock-sensitive unit tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp): cap cold schema validation workers at the CPUs the proxy can use Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: inject the cache clock and wait on lock contention instead of 1s deadlines Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(mcp): wait on listing events instead of 1s deadlines in the discovery drain test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo <mateo@berri.ai>
This commit is contained in:
parent
de0ed1557e
commit
bf56719f09
11 changed files with 166 additions and 44 deletions
|
|
@ -2,6 +2,7 @@ from __future__ import annotations
|
|||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from collections import deque
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
|
|
@ -55,6 +56,7 @@ _MAX_VALIDATION_DEPTH: Final = 64
|
|||
_MAX_VALIDATION_NODES: Final = 10_000
|
||||
_MAX_VALIDATION_CHARACTERS: Final = 1_048_576
|
||||
_VALIDATION_TIMEOUT_SECONDS: Final = 30
|
||||
_MAX_VALIDATION_WORKERS: Final = 4
|
||||
|
||||
|
||||
def _validation_nodes(value: JsonValue | Mapping[str, JsonValue]) -> Iterator[tuple[int, int]]:
|
||||
|
|
@ -101,6 +103,11 @@ def _validate_tool_arguments(schema: Mapping[str, JsonValue], arguments: Mapping
|
|||
return None
|
||||
|
||||
|
||||
def _schema_validation_worker_count() -> int:
|
||||
available: Final = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count() or 1
|
||||
return max(1, min(_MAX_VALIDATION_WORKERS, available))
|
||||
|
||||
|
||||
async def _tool_argument_validation_error(
|
||||
schema: Mapping[str, JsonValue],
|
||||
arguments: Mapping[str, JsonValue],
|
||||
|
|
@ -110,7 +117,7 @@ async def _tool_argument_validation_error(
|
|||
if limit_error is not None:
|
||||
return limit_error
|
||||
existing: Final = _SCHEMA_VALIDATION_LIMITER.get(None)
|
||||
limiter: Final = existing if existing is not None else anyio.CapacityLimiter(4)
|
||||
limiter: Final = existing if existing is not None else anyio.CapacityLimiter(_schema_validation_worker_count())
|
||||
if existing is None:
|
||||
_SCHEMA_VALIDATION_LIMITER.set(limiter)
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from __future__ import annotations
|
|||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from itertools import chain
|
||||
from typing import Any, Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
|
@ -837,7 +838,12 @@ async def test_the_prefetch_reads_cooldowns_again_once_the_read_interval_elapses
|
|||
@pytest.mark.asyncio
|
||||
async def test_the_prefetch_mget_carries_only_the_keys_whose_read_is_due(monkeypatch):
|
||||
first_time: Final = 1_000_000.0
|
||||
pinned: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc)
|
||||
monkeypatch.setattr(dual_cache_module.time, "time", lambda: first_time)
|
||||
monkeypatch.setattr(
|
||||
"litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime",
|
||||
lambda: pinned,
|
||||
)
|
||||
client = FakeClient(_lua_ok_replies)
|
||||
redis_cache = FakeRedisCache(client)
|
||||
router = _router(redis_cache, routing_strategy="usage-based-routing-v2")
|
||||
|
|
|
|||
|
|
@ -443,6 +443,42 @@ async def test_generic_api_callback_sumologic_uses_ndjson():
|
|||
assert len(mine) == 2, f"Expected this test's 2 calls as NDJSON lines, got {len(mine)} of {len(records)}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_api_callback_sends_queued_events_on_its_own(monkeypatch: pytest.MonkeyPatch):
|
||||
mock_post = AsyncMock()
|
||||
mock_post.return_value.status_code = 200
|
||||
mock_post.return_value.text = "OK"
|
||||
|
||||
test_endpoint = "https://example.com/api/logs"
|
||||
test_headers = {"Authorization": "Bearer test_token"}
|
||||
monkeypatch.setenv("GENERIC_LOGGER_ENDPOINT", test_endpoint)
|
||||
|
||||
generic_logger = GenericAPILogger(endpoint=test_endpoint, headers=test_headers, flush_interval=0.05)
|
||||
generic_logger.async_httpx_client.post = mock_post
|
||||
monkeypatch.setattr(litellm, "callbacks", [generic_logger])
|
||||
|
||||
for i in range(3):
|
||||
await litellm.acompletion(
|
||||
model="gpt-5.5",
|
||||
messages=[{"role": "user", "content": f"Hello, world! {i}"}],
|
||||
mock_response="hi",
|
||||
user="test_user",
|
||||
)
|
||||
|
||||
async def _poll() -> None:
|
||||
while sum(len(json.loads(call.kwargs["data"])) for call in mock_post.call_args_list) < 3:
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
await asyncio.wait_for(_poll(), timeout=10)
|
||||
|
||||
posted_bodies: Final = [json.loads(call.kwargs["data"]) for call in mock_post.call_args_list]
|
||||
for body in posted_bodies:
|
||||
assert isinstance(body, list), "Each posted body should be a JSON array"
|
||||
posted_items: Final = sum(posted_bodies, [])
|
||||
contents: Final = sorted(item["messages"][0]["content"] for item in posted_items)
|
||||
assert contents == [f"Hello, world! {i}" for i in range(3)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_api_callback_invalid_log_format():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1080,11 +1080,19 @@ def test_auth_check_and_project_id_make_one_round_trip_when_langfuse_is_down(
|
|||
[(500, 2), (503, 2), (429, 1), (404, 1)],
|
||||
ids=["http-500", "http-503", "http-429", "http-404"],
|
||||
)
|
||||
def test_cold_prompt_miss_never_sleeps_when_langfuse_is_down(status: int, round_trips: int):
|
||||
def test_cold_prompt_miss_never_sleeps_when_langfuse_is_down(
|
||||
status: int, round_trips: int, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""A cold ``get_prompt`` fetches inline on the event loop; with the generated client's default retries a
|
||||
429 carrying ``Retry-After: 30`` used to hold the loop for a minute. A 5xx gets the v2 client's one
|
||||
quick retry, a 429 or 4xx none."""
|
||||
requests: list[httpx.Request] = []
|
||||
sleeps: Final[list[float]] = []
|
||||
monkeypatch.setattr(
|
||||
langfuse_http_client,
|
||||
"time",
|
||||
SimpleNamespace(sleep=sleeps.append, time=time.time),
|
||||
)
|
||||
|
||||
def fail(request: httpx.Request) -> httpx.Response:
|
||||
requests.append(request)
|
||||
|
|
@ -1097,18 +1105,25 @@ def test_cold_prompt_miss_never_sleeps_when_langfuse_is_down(status: int, round_
|
|||
httpx_client=httpx.Client(transport=httpx.MockTransport(fail)),
|
||||
)
|
||||
|
||||
started = monotonic()
|
||||
with pytest.raises(LangfusePromptError) as caught:
|
||||
client.get_prompt("greeting")
|
||||
assert len(requests) == round_trips
|
||||
assert monotonic() - started < 0.5
|
||||
assert sleeps == [], "a cold prompt miss should not sleep for REST retries"
|
||||
assert caught.value.status_code == status
|
||||
|
||||
|
||||
@pytest.mark.parametrize("first_failure", [503, "connect-error"], ids=["http-503", "connect-error"])
|
||||
def test_one_transient_failure_on_a_cold_prompt_miss_does_not_fail_the_call(first_failure: int | str):
|
||||
def test_one_transient_failure_on_a_cold_prompt_miss_does_not_fail_the_call(
|
||||
first_failure: int | str, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""The v2 client retried a cold fetch once; a single Langfuse blip must not fail the LLM call."""
|
||||
requests: list[httpx.Request] = []
|
||||
sleeps: Final[list[float]] = []
|
||||
monkeypatch.setattr(
|
||||
langfuse_http_client,
|
||||
"time",
|
||||
SimpleNamespace(sleep=sleeps.append, time=time.time),
|
||||
)
|
||||
|
||||
def flaky(request: httpx.Request) -> httpx.Response:
|
||||
requests.append(request)
|
||||
|
|
@ -1125,10 +1140,9 @@ def test_one_transient_failure_on_a_cold_prompt_miss_does_not_fail_the_call(firs
|
|||
httpx_client=httpx.Client(transport=httpx.MockTransport(flaky)),
|
||||
)
|
||||
|
||||
started = monotonic()
|
||||
assert client.get_prompt("greeting").compile() == "hello"
|
||||
assert len(requests) == 2
|
||||
assert monotonic() - started < 0.5
|
||||
assert sleeps == [], "one transient failure should not sleep for REST retries"
|
||||
assert client.get_prompt("greeting").compile() == "hello", "the retried prompt is cached like any other"
|
||||
assert len(requests) == 2
|
||||
|
||||
|
|
|
|||
|
|
@ -4121,6 +4121,18 @@ async def test_truncated_jsonrpc_response_with_nested_method_skips_lock(
|
|||
)
|
||||
mcp_server._stateful_session_owners[session_id] = mcp_server._owner_fingerprint_for(owner_auth)
|
||||
|
||||
class _RecordingLock(asyncio.Lock):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.contended = asyncio.Event()
|
||||
|
||||
async def acquire(self) -> bool:
|
||||
if self.locked():
|
||||
self.contended.set()
|
||||
return await super().acquire()
|
||||
|
||||
session_lock: Final = _RecordingLock()
|
||||
mcp_server._stateful_session_locks[session_id] = session_lock
|
||||
gate = asyncio.Event()
|
||||
request_in_handle = asyncio.Event()
|
||||
response_handled = asyncio.Event()
|
||||
|
|
@ -4186,12 +4198,19 @@ async def test_truncated_jsonrpc_response_with_nested_method_skips_lock(
|
|||
),
|
||||
):
|
||||
req_task = asyncio.create_task(call(request_body))
|
||||
await asyncio.wait_for(request_in_handle.wait(), timeout=1.0)
|
||||
await request_in_handle.wait()
|
||||
|
||||
resp_task = asyncio.create_task(call(response_body))
|
||||
# Under a flat substring heuristic the response would acquire the
|
||||
# lock held by req_task and this wait would time out (deadlock).
|
||||
await asyncio.wait_for(response_handled.wait(), timeout=1.0)
|
||||
response_handled_task: Final = asyncio.create_task(response_handled.wait())
|
||||
contended_task: Final = asyncio.create_task(session_lock.contended.wait())
|
||||
_, pending = await asyncio.wait(
|
||||
(response_handled_task, contended_task),
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
for event_task in pending:
|
||||
event_task.cancel()
|
||||
await asyncio.gather(*pending, return_exceptions=True)
|
||||
assert response_handled.is_set(), "response POST contended for the request's session lock"
|
||||
|
||||
await resp_task
|
||||
assert not req_task.done()
|
||||
|
|
@ -4207,6 +4226,7 @@ async def test_truncated_jsonrpc_response_with_nested_method_skips_lock(
|
|||
finally:
|
||||
gate.set()
|
||||
mcp_server._stateful_session_auth_contexts.pop(session_id, None)
|
||||
mcp_server._stateful_session_auth_context_last_seen.pop(session_id, None)
|
||||
mcp_server._stateful_session_owners.pop(session_id, None)
|
||||
mcp_server._stateful_session_locks.pop(session_id, None)
|
||||
mcp_server._stateful_session_active_request_counts.pop(session_id, None)
|
||||
|
|
|
|||
|
|
@ -1427,7 +1427,12 @@ async def test_mcp_tool_search_leaves_the_listed_tools_slot_empty(monkeypatch: p
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cold_concurrent_schema_validation_accepts_valid_arguments() -> None:
|
||||
from litellm.proxy._experimental.mcp_server.tool_search import _tool_argument_validation_error
|
||||
import psutil
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.tool_search import (
|
||||
_MAX_VALIDATION_WORKERS,
|
||||
_tool_argument_validation_error,
|
||||
)
|
||||
|
||||
schema: Final = {
|
||||
"type": "object",
|
||||
|
|
@ -1435,6 +1440,7 @@ async def test_cold_concurrent_schema_validation_accepts_valid_arguments() -> No
|
|||
"properties": {"amount": {"$ref": "#/$defs/amount"}},
|
||||
}
|
||||
original_affinity: Final = os.sched_getaffinity(0) if sys.platform == "linux" else None
|
||||
existing_children: Final = frozenset(child.pid for child in psutil.Process().children())
|
||||
try:
|
||||
if original_affinity is not None:
|
||||
os.sched_setaffinity(0, {min(original_affinity)})
|
||||
|
|
@ -1443,6 +1449,14 @@ async def test_cold_concurrent_schema_validation_accepts_valid_arguments() -> No
|
|||
*(_tool_argument_validation_error(schema, {"amount": 0.75}, timeout_seconds=90.0) for _ in range(8))
|
||||
)
|
||||
assert results == [None] * 8, "Cold and warm workers must accept valid concurrent tool arguments"
|
||||
workers: Final = tuple(
|
||||
child
|
||||
for child in psutil.Process().children()
|
||||
if child.pid not in existing_children and "anyio.to_process" in child.cmdline()
|
||||
)
|
||||
assert 1 <= len(workers) <= (1 if original_affinity is not None else _MAX_VALIDATION_WORKERS), (
|
||||
"Validation must not start more cold workers than the CPUs it can run on"
|
||||
)
|
||||
finally:
|
||||
if original_affinity is not None:
|
||||
os.sched_setaffinity(0, original_affinity)
|
||||
|
|
|
|||
|
|
@ -1287,17 +1287,17 @@ async def test_discovery_concurrent_listings_drain_on_failure_and_cancellation(o
|
|||
):
|
||||
task = asyncio.create_task(GatewayOperations().execute(DiscoverRequest(), prepare_context(UserAPIKeyAuth(user_id="scoped"))))
|
||||
try:
|
||||
await asyncio.wait_for(asyncio.gather(*(event.wait() for event in ready)), 1)
|
||||
await asyncio.gather(*(event.wait() for event in ready))
|
||||
if outcome == "cancel":
|
||||
task.cancel()
|
||||
else:
|
||||
release.set()
|
||||
if outcome == "success":
|
||||
result = await asyncio.wait_for(task, 1)
|
||||
result = await task
|
||||
assert result.capabilities.model_dump(exclude_none=True) == {}
|
||||
else:
|
||||
with pytest.raises(asyncio.CancelledError if outcome == "cancel" else ValueError):
|
||||
await asyncio.wait_for(task, 1)
|
||||
await task
|
||||
assert all(event.is_set() for event in closed)
|
||||
assert tools.call_args.kwargs["log_list_tools_to_spendlogs"] is False
|
||||
finally:
|
||||
|
|
|
|||
|
|
@ -499,12 +499,14 @@ async def test_native_count_finishing_after_quarter_worker_budget_keeps_plan_and
|
|||
from litellm.litellm_core_utils import logging_worker
|
||||
from litellm.litellm_core_utils.logging_worker import LoggingWorker
|
||||
|
||||
release: Final = asyncio.Event()
|
||||
counted: Final = asyncio.Event()
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
real_time: Final = loop.time
|
||||
monkeypatch.setattr(loop, "time", lambda: real_time() + (2.05 if counted.is_set() else 0.0))
|
||||
|
||||
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.05, release.set)
|
||||
await release.wait()
|
||||
counted.set()
|
||||
await asyncio.sleep(0.001)
|
||||
return await _count(model, api_key, body)
|
||||
|
||||
worker: Final = LoggingWorker(timeout=8.0)
|
||||
|
|
@ -519,7 +521,6 @@ async def test_native_count_finishing_after_quarter_worker_budget_keeps_plan_and
|
|||
assert observed.plan is not None and observed.plan.breakpoints[0].prefix_tokens == 5000
|
||||
assert payload["response_cost"] is not None and worker._timeout_total == 0
|
||||
finally:
|
||||
release.set()
|
||||
await worker.stop()
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -48,15 +48,15 @@ async def test_fork_lookups_overlap_with_a_bounded_number_of_requests() -> None:
|
|||
source: Final = GitLab(ROISettings(source_provider="gitlab"), httpx.MockTransport(respond))
|
||||
pending: Final = asyncio.create_task(source.pulls("group/repo", date(2026, 9, 1), date(2026, 9, 30)))
|
||||
try:
|
||||
first_wave: Final = tuple([await asyncio.wait_for(started.get(), timeout=1) for _ in range(8)])
|
||||
first_wave: Final = tuple([await started.get() for _ in range(8)])
|
||||
assert len(set(first_wave)) == 8
|
||||
assert started.empty()
|
||||
release[first_wave[0] - 2].set()
|
||||
next_id: Final = await asyncio.wait_for(started.get(), timeout=1)
|
||||
next_id: Final = await started.get()
|
||||
assert next_id not in first_wave
|
||||
for event in release:
|
||||
event.set()
|
||||
pulls: Final = await asyncio.wait_for(pending, timeout=1)
|
||||
pulls: Final = await pending
|
||||
assert tuple(pull.head.repo.full_name if pull.head and pull.head.repo else None for pull in pulls) == tuple(
|
||||
None if source_id == 3 else f"fork-{source_id}/repo" for source_id in source_ids
|
||||
)
|
||||
|
|
|
|||
|
|
@ -27,6 +27,17 @@ from litellm.proxy.utils import hash_token
|
|||
from litellm.proxy.db.log_db_metrics import record_db_io
|
||||
|
||||
|
||||
class _FakeClock:
|
||||
def __init__(self) -> None:
|
||||
self.now = 1000.0
|
||||
|
||||
def advance(self, seconds: float) -> None:
|
||||
self.now += seconds
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
|
||||
def _digest_row(digest: str, key_alias: str | None, team_id: str | None, user_id: str | None) -> dict[str, str | None]:
|
||||
return {"digest": digest, "key_alias": key_alias, "team_id": team_id, "user_id": user_id}
|
||||
|
||||
|
|
@ -422,34 +433,44 @@ async def test_recover_key_metadata_from_spend_logs_rescans_when_the_window_chan
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recover_key_metadata_from_spend_logs_retries_a_failed_query_only_after_the_miss_ttl():
|
||||
digest = hash_token("cli-session-retry")
|
||||
window = (datetime(2026, 9, 7), datetime(2026, 9, 10))
|
||||
cache = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL)
|
||||
mock_prisma = MagicMock()
|
||||
query_raw = _spend_log_transaction(mock_prisma, AsyncMock(side_effect=PrismaError("statement timeout")))
|
||||
started = time.time()
|
||||
assert await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) == {}
|
||||
query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([_spend_log_row(digest, "back-online", None, None)]))
|
||||
async def test_recover_key_metadata_from_spend_logs_retries_a_failed_query_only_after_the_miss_ttl() -> None:
|
||||
digest: Final = hash_token("cli-session-retry")
|
||||
window: Final = (datetime(2026, 9, 7), datetime(2026, 9, 10))
|
||||
clock: Final = _FakeClock()
|
||||
cache: Final = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL, clock=clock)
|
||||
mock_prisma: Final = MagicMock()
|
||||
query_raw: Final = _spend_log_transaction(
|
||||
mock_prisma,
|
||||
AsyncMock(
|
||||
side_effect=[
|
||||
PrismaError("statement timeout"),
|
||||
[_spend_log_row(digest, "back-online", None, None)],
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
assert await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) == {}
|
||||
query_raw.assert_not_awaited()
|
||||
miss_key = next(key for key in cache.ttl_dict if digest in key and not key.endswith(":missed-before"))
|
||||
assert cache.ttl_dict[miss_key] - started <= SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1
|
||||
cache.ttl_dict[miss_key] = time.time() - 1
|
||||
miss_key: Final = next(key for key in cache.ttl_dict if digest in key and not key.endswith(":missed-before"))
|
||||
assert cache.ttl_dict[miss_key] == clock.now + SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL
|
||||
|
||||
result = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache)
|
||||
assert await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) == {}
|
||||
assert query_raw.await_count == 1
|
||||
clock.advance(SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1)
|
||||
|
||||
result: Final = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache)
|
||||
|
||||
assert result[digest]["key_alias"] == "back-online"
|
||||
assert query_raw.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recover_key_metadata_from_spend_logs_never_caches_repeated_query_failures_as_long_as_a_hit():
|
||||
async def test_recover_key_metadata_from_spend_logs_never_caches_repeated_query_failures_as_long_as_a_hit() -> None:
|
||||
digest: Final = hash_token("cli-session-repeated-failure")
|
||||
window: Final = (datetime(2026, 9, 7), datetime(2026, 9, 10))
|
||||
cache: Final = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL)
|
||||
clock: Final = _FakeClock()
|
||||
cache: Final = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL, clock=clock)
|
||||
mock_prisma: Final = MagicMock()
|
||||
_spend_log_transaction(
|
||||
query_raw: Final = _spend_log_transaction(
|
||||
mock_prisma,
|
||||
AsyncMock(
|
||||
side_effect=[
|
||||
|
|
@ -461,17 +482,19 @@ async def test_recover_key_metadata_from_spend_logs_never_caches_repeated_query_
|
|||
)
|
||||
await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache)
|
||||
miss_key: Final = next(key for key in cache.ttl_dict if digest in key and not key.endswith(":missed-before"))
|
||||
cache.ttl_dict[miss_key] = time.time() - 1
|
||||
second_query_started: Final = time.time()
|
||||
assert cache.ttl_dict[miss_key] == clock.now + SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL
|
||||
clock.advance(SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1)
|
||||
|
||||
await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache)
|
||||
|
||||
assert cache.ttl_dict[miss_key] - second_query_started <= SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1
|
||||
cache.ttl_dict[miss_key] = time.time() - 1
|
||||
assert query_raw.await_count == 2
|
||||
assert cache.ttl_dict[miss_key] == clock.now + SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL
|
||||
clock.advance(SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1)
|
||||
|
||||
result: Final = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache)
|
||||
|
||||
assert result[digest]["key_alias"] == "back-online"
|
||||
assert query_raw.await_count == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -1071,6 +1071,7 @@ def _shared_redis_stub(store: dict) -> MagicMock:
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("router_minute_pinned")
|
||||
async def test_headers_on_fresh_worker_reflect_shared_redis_usage():
|
||||
store: dict = {}
|
||||
worker_a = _rpm_tpm_router("lit-3058-workers")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue