mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-10 08:28:25 +00:00
Merge 1343f38da9 into f4d1f45d87
This commit is contained in:
yucheng-berri
GitHub
commit
f87c792763
6 changed files with 400 additions and 43 deletions
|
|
@ -453,6 +453,16 @@ def is_redis_timeout_failure(exc: BaseException) -> bool:
|
|||
return any(isinstance(link, timeout_types) for link in _explicit_causes(exc))
|
||||
|
||||
|
||||
def is_redis_scripting_unavailable(exc: BaseException) -> bool:
|
||||
"""True when Redis refused to run a Lua script for this client: an ACL without
|
||||
``@scripting``, or ``EVAL`` renamed away."""
|
||||
from redis.exceptions import NoPermissionError, ResponseError
|
||||
|
||||
if isinstance(exc, NoPermissionError):
|
||||
return "permissions to run" in str(exc)
|
||||
return isinstance(exc, ResponseError) and "unknown command" in str(exc).lower()
|
||||
|
||||
|
||||
def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseException) -> None:
|
||||
"""Record a Redis failure that the calling method is about to swallow.
|
||||
|
||||
|
|
@ -1555,9 +1565,36 @@ class RedisCache(BaseCache):
|
|||
"return ARGV[1] end "
|
||||
"return cur"
|
||||
)
|
||||
result = cast(
|
||||
"str | bytes | int | float | None",
|
||||
await _redis_client.eval(lua, 1, key, str(value), str(int(_used_ttl or 0))),
|
||||
start_time: Final = time.time()
|
||||
try:
|
||||
result = cast(
|
||||
"str | bytes | int | float | None",
|
||||
await _redis_client.eval(lua, 1, key, str(value), str(int(_used_ttl or 0))),
|
||||
)
|
||||
except Exception as e:
|
||||
failed_at: Final = time.time()
|
||||
asyncio.create_task(
|
||||
self.service_logger_obj.async_service_failure_hook( # pyright: ignore[reportUnknownMemberType] # bare dict in sig
|
||||
service=ServiceTypes.REDIS,
|
||||
duration=failed_at - start_time,
|
||||
error=e,
|
||||
call_type="async_set_max",
|
||||
caller=_get_call_stack_info(),
|
||||
start_time=start_time,
|
||||
end_time=failed_at,
|
||||
)
|
||||
)
|
||||
raise
|
||||
end_time: Final = time.time()
|
||||
asyncio.create_task(
|
||||
self.service_logger_obj.async_service_success_hook( # pyright: ignore[reportUnknownMemberType] # bare dict in sig
|
||||
service=ServiceTypes.REDIS,
|
||||
duration=end_time - start_time,
|
||||
call_type="async_set_max",
|
||||
caller=_get_call_stack_info(),
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
)
|
||||
if result is None:
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -215,6 +215,7 @@ _SERVICE_VERB_BY_CALL_TYPE: Final[dict[str, str]] = {
|
|||
"async_increment": "incr",
|
||||
"async_increment_pipeline": "incr",
|
||||
"async_increment_if_present": "incr",
|
||||
"async_set_max": "set",
|
||||
"delete_cache": "delete",
|
||||
"async_delete_cache": "delete",
|
||||
"async_rpush": "rpush",
|
||||
|
|
|
|||
|
|
@ -323,7 +323,11 @@ from litellm.caching.redis_batch import (
|
|||
active_post_call_redis_batch,
|
||||
drain_post_call_redis_batches,
|
||||
)
|
||||
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, is_redis_timeout_failure
|
||||
from litellm.caching.redis_cache import (
|
||||
RedisCircuitBreakerOpenError,
|
||||
is_redis_scripting_unavailable,
|
||||
is_redis_timeout_failure,
|
||||
)
|
||||
from litellm.caching.redis_cluster_cache import RedisClusterCache
|
||||
from litellm.constants import (
|
||||
_REALTIME_BODY_CACHE_SIZE,
|
||||
|
|
@ -3349,15 +3353,7 @@ async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None
|
|||
in-memory copy is guarded by a read-compare-write with no await in between,
|
||||
so it is atomic within the worker.
|
||||
"""
|
||||
cached: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
|
||||
needs_update = True
|
||||
if cached is not None:
|
||||
try:
|
||||
needs_update = float(cached) < db_spend
|
||||
except (TypeError, ValueError):
|
||||
needs_update = True
|
||||
if needs_update:
|
||||
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=db_spend)
|
||||
_lift_in_memory_spend_counter(counter_key=counter_key, floor=db_spend)
|
||||
if spend_counter_cache.redis_cache is not None:
|
||||
forget_spend_counter(counter_key)
|
||||
try:
|
||||
|
|
@ -3373,6 +3369,18 @@ async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None
|
|||
record_spend_counter_value(counter_key, repaired)
|
||||
|
||||
|
||||
def _lift_in_memory_spend_counter(counter_key: str, floor: float) -> None:
|
||||
cached: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
|
||||
needs_update = True
|
||||
if cached is not None:
|
||||
try:
|
||||
needs_update = float(cached) < floor
|
||||
except (TypeError, ValueError):
|
||||
needs_update = True
|
||||
if needs_update:
|
||||
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=floor)
|
||||
|
||||
|
||||
async def reseed_spend_counter_from_db(counter_key: str) -> bool:
|
||||
"""Recover a counter that the reservation reconcile found in an inconsistent
|
||||
state (missing, or where applying the reconcile delta would drive it
|
||||
|
|
@ -3994,11 +4002,8 @@ async def _prepare_spend_counter_increment(
|
|||
2. If not found, reseed from the DB via `SpendCounterReseed.coalesced`.
|
||||
Falls back to the cached object's `.spend` via user_api_key_cache
|
||||
only if prisma is unavailable, since that value can lag the flusher.
|
||||
3. Seed counter via async_increment_cache (not async_set_cache) to avoid a
|
||||
check-then-set race: if two pods cold-start simultaneously, both may see
|
||||
the counter as absent and seed it. Using increment means the worst case
|
||||
is over-counting (conservative, blocks slightly early) rather than
|
||||
under-counting (would allow overspend).
|
||||
3. Lift the counter to at least that spend rather than adding it, so
|
||||
requests and pods that cold-start it together count it once.
|
||||
4. Increment is returned for the caller to apply via pipeline
|
||||
"""
|
||||
with service_target(SPEND_COUNTERS_TARGET):
|
||||
|
|
@ -4103,7 +4108,7 @@ async def ensure_spend_counter_initialized(
|
|||
# DB unavailable - fall back to in-process cache (may be stale).
|
||||
base_spend: Final = await _get_source_cache_base_spend(source_cache_key=source_cache_key)
|
||||
if base_spend > 0:
|
||||
await increment_spend_counter_cache(counter_key=counter_key, increment=base_spend)
|
||||
await _lift_spend_counter_cache(counter_key=counter_key, floor=base_spend)
|
||||
|
||||
|
||||
_ensure_spend_counter_initialized: Final = ensure_spend_counter_initialized
|
||||
|
|
@ -4232,6 +4237,24 @@ async def increment_spend_counter_cache(counter_key: str, increment: float) -> f
|
|||
_increment_spend_counter_cache: Final = increment_spend_counter_cache
|
||||
|
||||
|
||||
@with_service_target(SPEND_COUNTERS_TARGET)
|
||||
async def _lift_spend_counter_cache(counter_key: str, floor: float) -> None:
|
||||
if spend_counter_cache.redis_cache is not None:
|
||||
try:
|
||||
lifted: Final = await spend_counter_cache.redis_cache.async_set_max(key=counter_key, value=floor)
|
||||
except Exception as e:
|
||||
if is_redis_scripting_unavailable(e):
|
||||
await increment_spend_counter_cache(counter_key=counter_key, increment=floor)
|
||||
return
|
||||
await invalidate_spend_counter(counter_key=counter_key)
|
||||
raise
|
||||
if lifted is not None:
|
||||
_lift_in_memory_spend_counter(counter_key=counter_key, floor=lifted)
|
||||
record_spend_counter_value(counter_key, lifted)
|
||||
return
|
||||
_lift_in_memory_spend_counter(counter_key=counter_key, floor=floor)
|
||||
|
||||
|
||||
@with_service_target(SPEND_COUNTERS_TARGET)
|
||||
async def invalidate_spend_counter(counter_key: str) -> None:
|
||||
forget_spend_counter(counter_key)
|
||||
|
|
|
|||
|
|
@ -1,17 +1,20 @@
|
|||
import json
|
||||
import os
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Callable, Mapping
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.database import read_rows, write_rows
|
||||
from integration._support.upstream import delete_scenario, register_scenario
|
||||
from integration.cost_calculation.cost_tracking_case import JsonResponse
|
||||
from pydantic import BaseModel, JsonValue
|
||||
from redis import Redis
|
||||
|
||||
PROVIDER_AUTHORIZATION: Final = "Bearer integration-provider-key"
|
||||
CANNED_CONTENT: Final = "Hello! This is a mock response from the fake OpenAI endpoint."
|
||||
|
|
@ -639,3 +642,135 @@ def test_streamed_call_is_refused_as_json_once_over_budget(gateway: Gateway) ->
|
|||
assert "data:" not in denied.text, denied.text
|
||||
_assert_budget_refusal(denied, end_user, 0.05)
|
||||
assert _observed_requests(upstream) == []
|
||||
|
||||
|
||||
_COLD_BURST: Final = 40
|
||||
_COLD_SEED_SPEND: Final = 6.0
|
||||
_COLD_BUDGET: Final = 10.0
|
||||
_CHAT_COST: Final = 20 * 0.001 + 20 * 0.002
|
||||
|
||||
|
||||
def _serves_model(gateway: Gateway, model: str) -> bool:
|
||||
response: Final = gateway.request("GET", "/model/info")
|
||||
return response.status_code == 200 and any(
|
||||
object_value(entry)["model_name"] == model for entry in object_value(response.json())["data"]
|
||||
)
|
||||
|
||||
|
||||
def _counter_cache() -> Redis:
|
||||
return Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]))
|
||||
|
||||
|
||||
def _counter(cache: Redis, end_user: str) -> float | None:
|
||||
raw: Final = cache.get(f"spend:end_user:{end_user}")
|
||||
return None if raw is None else float(raw)
|
||||
|
||||
|
||||
def _lose_counter(cache: Redis, end_user: str) -> None:
|
||||
cache.delete(f"spend:end_user:{end_user}")
|
||||
|
||||
|
||||
def _set_end_user_row_spend(end_user: str, spend: float) -> None:
|
||||
write_rows('UPDATE "LiteLLM_EndUserTable" SET spend = %s WHERE user_id = %s', (str(spend), end_user))
|
||||
|
||||
|
||||
def _await_row_spend(end_user: str, spend: float) -> None:
|
||||
eventually(
|
||||
lambda: _end_user_row(end_user),
|
||||
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) == pytest.approx(spend),
|
||||
seconds=70,
|
||||
)
|
||||
|
||||
|
||||
def _refused_spend(denied: httpx.Response, end_user: str) -> float:
|
||||
error: Final = _denial_error(denied)
|
||||
assert error["type"] == "budget_exceeded", denied.text
|
||||
matched: Final = _REFUSAL_MESSAGE.match(string_value(error["message"]))
|
||||
assert matched is not None, denied.text
|
||||
assert matched.group("user") == end_user
|
||||
return float(matched.group("spend"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("spelling", ("user", "responses_user", "messages_metadata_user_id"))
|
||||
def test_overlapping_cold_starts_count_the_end_user_spend_once(gateway: Gateway, peer: Gateway, spelling: str) -> None:
|
||||
chosen: Final = SPELLINGS[spelling]
|
||||
with (
|
||||
gateway.scenario() as scenario,
|
||||
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
||||
_counter_cache() as cache,
|
||||
):
|
||||
deployment: Final = _deployment(scenario, chosen)
|
||||
eventually(lambda: _serves_model(peer, deployment.model), bool, seconds=30)
|
||||
key: Final = scenario.key(models=[deployment.model])
|
||||
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
||||
_new_customer(gateway, end_user, max_budget=_COLD_BUDGET, spend=_COLD_SEED_SPEND)
|
||||
scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
||||
assert _customer_info(gateway, end_user).spend == _COLD_SEED_SPEND
|
||||
assert _counter(cache, end_user) is None
|
||||
upstream.delete("/__observations").raise_for_status()
|
||||
|
||||
proxies: Final = (gateway, peer)
|
||||
with ThreadPoolExecutor(max_workers=_COLD_BURST) as pool:
|
||||
burst: Final = tuple(
|
||||
pool.map(
|
||||
lambda index: _send(proxies[index % 2], chosen, deployment, key, end_user, f"cold {index}"),
|
||||
range(_COLD_BURST),
|
||||
)
|
||||
)
|
||||
refused: Final = [response.text for response in burst if response.status_code != 200]
|
||||
assert refused == [], refused
|
||||
follow_up: Final = _send(gateway, chosen, deployment, key, end_user, "cold follow-up")
|
||||
_assert_served(chosen, follow_up)
|
||||
assert len(_observed_requests(upstream)) == _COLD_BURST + 1
|
||||
|
||||
expected_spend: Final = _COLD_SEED_SPEND + (_COLD_BURST + 1) * _CHAT_COST
|
||||
_await_row_spend(end_user, expected_spend)
|
||||
assert _customer_info(gateway, end_user).spend == pytest.approx(expected_spend)
|
||||
counted: Final = eventually(
|
||||
lambda: _counter(cache, end_user),
|
||||
lambda value: value is not None and value == pytest.approx(expected_spend),
|
||||
seconds=30,
|
||||
return_last_on_timeout=True,
|
||||
)
|
||||
assert counted == pytest.approx(expected_spend)
|
||||
|
||||
|
||||
def test_a_lost_counter_keeps_spend_this_proxy_has_not_written_yet(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario, _counter_cache() as cache:
|
||||
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.25)
|
||||
key: Final = scenario.key(models=[model])
|
||||
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
||||
call_cost: Final = 20 * 0.001 + 20 * 0.25
|
||||
_new_customer(gateway, end_user, max_budget=_COLD_BUDGET, spend=_COLD_SEED_SPEND)
|
||||
scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
||||
served: Final = _chat(gateway, model, key, end_user)
|
||||
assert served.status_code == 200, served.text
|
||||
_await_row_spend(end_user, _COLD_SEED_SPEND + call_cost)
|
||||
_set_end_user_row_spend(end_user, _COLD_SEED_SPEND)
|
||||
_lose_counter(cache, end_user)
|
||||
|
||||
denied: Final = _chat(gateway, model, key, end_user)
|
||||
|
||||
assert _refused_spend(denied, end_user) == pytest.approx(_COLD_SEED_SPEND + call_cost)
|
||||
|
||||
|
||||
def test_a_lost_counter_keeps_spend_another_proxy_has_not_written_yet(gateway: Gateway, peer: Gateway) -> None:
|
||||
with gateway.scenario() as scenario, _counter_cache() as cache:
|
||||
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.149)
|
||||
eventually(lambda: _serves_model(peer, model), bool, seconds=30)
|
||||
key: Final = scenario.key(models=[model])
|
||||
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
||||
call_cost: Final = 20 * 0.001 + 20 * 0.149
|
||||
_new_customer(gateway, end_user, max_budget=_COLD_BUDGET, spend=_COLD_SEED_SPEND)
|
||||
scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
||||
first: Final = _chat(peer, model, key, end_user)
|
||||
assert first.status_code == 200, first.text
|
||||
_await_row_spend(end_user, _COLD_SEED_SPEND + call_cost)
|
||||
_set_end_user_row_spend(end_user, _COLD_SEED_SPEND)
|
||||
_lose_counter(cache, end_user)
|
||||
second: Final = _chat(peer, model, key, end_user)
|
||||
assert second.status_code == 200, second.text
|
||||
|
||||
denied: Final = _chat(gateway, model, key, end_user)
|
||||
|
||||
assert _refused_spend(denied, end_user) == pytest.approx(_COLD_SEED_SPEND + 2 * call_cost)
|
||||
|
|
|
|||
|
|
@ -1697,3 +1697,43 @@ async def test_async_set_cache_pipeline_does_not_format_cached_values_for_loggin
|
|||
|
||||
assert _FormatCountingStr.format_calls == 0
|
||||
assert pipe.sets == [("k1", '"embedding"', timedelta(seconds=60))]
|
||||
|
||||
|
||||
def _cache_reporting_to(service_logger: ServiceLogging) -> RedisCache:
|
||||
with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point
|
||||
"litellm._redis.get_redis_client", return_value=MagicMock()
|
||||
):
|
||||
return RedisCache(host="127.0.0.1", port=6379, service_logger_obj=service_logger)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_set_max_reports_a_failed_script_to_the_redis_failure_metric(redis_no_ping: None):
|
||||
"""The cold spend-counter seed runs through this script, so its Redis errors have to keep reaching
|
||||
litellm_redis_failed_requests_total the way they did when the seed was an INCRBYFLOAT."""
|
||||
service_logger: Final = ServiceLogging(mock_testing=True)
|
||||
cache: Final = _cache_reporting_to(service_logger)
|
||||
client: Final = AsyncMock()
|
||||
client.eval.side_effect = OSError("redis unavailable")
|
||||
tasks_before: Final = asyncio.all_tasks()
|
||||
|
||||
with patch.object(cache, "init_async_client", return_value=client):
|
||||
with pytest.raises(OSError, match="redis unavailable"):
|
||||
await cache.async_set_max(key="spend:end_user:customer-1", value=6.0)
|
||||
await asyncio.wait_for(asyncio.gather(*(asyncio.all_tasks() - tasks_before)), timeout=5)
|
||||
|
||||
assert (service_logger.mock_testing_async_success_hook, service_logger.mock_testing_async_failure_hook) == (0, 1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_set_max_reports_a_completed_script_to_the_redis_latency_metric(redis_no_ping: None):
|
||||
service_logger: Final = ServiceLogging(mock_testing=True)
|
||||
cache: Final = _cache_reporting_to(service_logger)
|
||||
client: Final = AsyncMock()
|
||||
client.eval.return_value = b"6.0"
|
||||
tasks_before: Final = asyncio.all_tasks()
|
||||
|
||||
with patch.object(cache, "init_async_client", return_value=client):
|
||||
assert await cache.async_set_max(key="spend:end_user:customer-1", value=6.0) == 6.0
|
||||
await asyncio.wait_for(asyncio.gather(*(asyncio.all_tasks() - tasks_before)), timeout=5)
|
||||
|
||||
assert (service_logger.mock_testing_async_success_hook, service_logger.mock_testing_async_failure_hook) == (1, 0)
|
||||
|
|
|
|||
|
|
@ -33,11 +33,13 @@ from opentelemetry.sdk.trace import TracerProvider
|
|||
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
|
||||
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
|
||||
from opentelemetry.trace import StatusCode
|
||||
from redis.exceptions import NoPermissionError, ResponseError
|
||||
|
||||
import litellm
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm._service_logger import ServiceLogging
|
||||
from litellm._internal_context import current_service_target
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.caching.redis_cache import RedisCache, RedisCircuitBreaker
|
||||
from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY
|
||||
|
|
@ -1329,32 +1331,151 @@ async def test_ensure_spend_counter_initialized_warm_skips_reseed_and_source(
|
|||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ensure_spend_counter_initialized_cold_seeds_from_source_cache(
|
||||
monkeypatch,
|
||||
):
|
||||
fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=7.0)
|
||||
fake_user_cache = _make_user_api_key_cache(get_value={"spend": 7.0})
|
||||
monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
|
||||
monkeypatch.setattr(ps, "prisma_client", None)
|
||||
monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None))
|
||||
class _SharedRedisCounters:
|
||||
"""The Redis counter commands a cold start uses, yielding first like a network call so overlapping starts
|
||||
interleave the way they do across workers and pods."""
|
||||
|
||||
await ps._ensure_spend_counter_initialized(
|
||||
counter_key="spend:user:u",
|
||||
source_cache_key="u",
|
||||
def __init__(self) -> None:
|
||||
self.values: Final[dict[str, float]] = {} # mutable-ok: stands in for the Redis keyspace
|
||||
|
||||
async def async_get_cache(self, key: str, **_: object) -> float | None:
|
||||
await asyncio.sleep(0)
|
||||
return self.values.get(key)
|
||||
|
||||
async def async_increment(self, key: str, value: float, **_: object) -> float:
|
||||
await asyncio.sleep(0)
|
||||
self.values[key] = self.values.get(key, 0.0) + value
|
||||
return self.values[key]
|
||||
|
||||
async def async_set_max(self, key: str, value: float, **_: object) -> float:
|
||||
await asyncio.sleep(0)
|
||||
self.values[key] = max(self.values.get(key, value), value)
|
||||
return self.values[key]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("with_redis", (True, False))
|
||||
@pytest.mark.parametrize("counter_key", ("spend:end_user:customer-1", "spend:tag:tag-1"))
|
||||
async def test_overlapping_cold_starts_count_the_cached_spend_once(
|
||||
monkeypatch: pytest.MonkeyPatch, with_redis: bool, counter_key: str
|
||||
) -> None:
|
||||
cached_spend: Final = {"first-start": 2.5, "second-start": 6.0, "third-start": 6.0, "fourth-start": 2.5}
|
||||
|
||||
async def _cached_object(key: str, **_: object) -> dict[str, float]:
|
||||
await asyncio.sleep(0)
|
||||
return {"spend": cached_spend[key]}
|
||||
|
||||
redis: Final = _SharedRedisCounters()
|
||||
counters: Final = DualCache(
|
||||
in_memory_cache=InMemoryCache(), redis_cache=cast(RedisCache, redis) if with_redis else None
|
||||
)
|
||||
objects: Final = MagicMock()
|
||||
objects.async_get_cache = AsyncMock(side_effect=_cached_object)
|
||||
monkeypatch.setattr(ps, "spend_counter_cache", counters)
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", objects)
|
||||
monkeypatch.setattr(ps, "prisma_client", None)
|
||||
|
||||
await asyncio.gather(
|
||||
*(
|
||||
ps.ensure_spend_counter_initialized(counter_key=counter_key, source_cache_key=start)
|
||||
for start in cached_spend
|
||||
)
|
||||
)
|
||||
|
||||
observed = {
|
||||
"source_cache_called": fake_user_cache.async_get_cache.called,
|
||||
"seed_increment_called": fake_cache.redis_cache.async_increment.called,
|
||||
"warm_check_done": fake_cache.redis_cache.async_get_cache.called,
|
||||
}
|
||||
assert normalize(observed) == {
|
||||
"source_cache_called": True,
|
||||
"seed_increment_called": True,
|
||||
"warm_check_done": True,
|
||||
}
|
||||
seeded: Final = redis.values.get(counter_key) if with_redis else counters.in_memory_cache.get_cache(key=counter_key)
|
||||
assert seeded == 6.0
|
||||
|
||||
|
||||
class _RedisThatFailsTheLift:
|
||||
async def async_get_cache(self, key: str, **_: object) -> None:
|
||||
return None
|
||||
|
||||
async def async_set_max(self, key: str, value: float, **_: object) -> float:
|
||||
raise ConnectionError("redis went away")
|
||||
|
||||
async def async_delete_cache(self, key: str, **_: object) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_redis_failure_while_lifting_a_cold_counter_drops_the_local_copy_and_raises(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
counter_key: Final = "spend:end_user:customer-1"
|
||||
counters: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=cast(RedisCache, _RedisThatFailsTheLift()))
|
||||
counters.in_memory_cache.set_cache(key=counter_key, value=2.0)
|
||||
objects: Final = MagicMock()
|
||||
objects.async_get_cache = AsyncMock(return_value={"spend": 6.0})
|
||||
monkeypatch.setattr(ps, "spend_counter_cache", counters)
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", objects)
|
||||
monkeypatch.setattr(ps, "prisma_client", None)
|
||||
|
||||
with pytest.raises(ConnectionError):
|
||||
await ps.ensure_spend_counter_initialized(counter_key=counter_key, source_cache_key="customer-1")
|
||||
|
||||
assert counters.in_memory_cache.get_cache(key=counter_key) is None
|
||||
|
||||
|
||||
class _FakeRedisRefusingScripts(FakeAsyncRedis):
|
||||
def __init__(self, server: FakeServer, refusal: ResponseError) -> None:
|
||||
super().__init__(server=server)
|
||||
self._refusal: Final = refusal
|
||||
|
||||
async def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> bytes | None:
|
||||
raise self._refusal
|
||||
|
||||
|
||||
def _cold_start_on(monkeypatch: pytest.MonkeyPatch, redis: FakeAsyncRedis, cached_spend: float) -> None:
|
||||
objects: Final = MagicMock()
|
||||
objects.async_get_cache = AsyncMock(return_value={"spend": cached_spend})
|
||||
monkeypatch.setattr(
|
||||
ps, "spend_counter_cache", DualCache(in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnClient(redis))
|
||||
)
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", objects)
|
||||
monkeypatch.setattr(ps, "prisma_client", None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"refusal",
|
||||
(
|
||||
NoPermissionError("User noscript has no permissions to run the 'eval' command"),
|
||||
ResponseError("unknown command 'eval', with args beginning with: "),
|
||||
),
|
||||
ids=("acl-without-scripting", "eval-renamed-away"),
|
||||
)
|
||||
async def test_a_cold_counter_is_still_seeded_on_a_redis_that_refuses_scripts(
|
||||
monkeypatch: pytest.MonkeyPatch, refusal: ResponseError
|
||||
) -> None:
|
||||
redis: Final = _FakeRedisRefusingScripts(FakeServer(), refusal)
|
||||
_cold_start_on(monkeypatch, redis, cached_spend=6.0)
|
||||
|
||||
await ps.ensure_spend_counter_initialized(counter_key="spend:end_user:customer-1", source_cache_key="customer-1")
|
||||
|
||||
assert await redis.get("spend:end_user:customer-1") == b"6"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"refusal",
|
||||
(
|
||||
NoPermissionError("No permissions to access a key"),
|
||||
ResponseError("WRONGTYPE Operation against a key holding the wrong kind of value"),
|
||||
),
|
||||
ids=("key-outside-the-acl", "wrong-type"),
|
||||
)
|
||||
async def test_a_lift_refused_for_another_reason_is_raised_instead_of_seeded_by_adding(
|
||||
monkeypatch: pytest.MonkeyPatch, refusal: ResponseError
|
||||
) -> None:
|
||||
redis: Final = _FakeRedisRefusingScripts(FakeServer(), refusal)
|
||||
_cold_start_on(monkeypatch, redis, cached_spend=6.0)
|
||||
|
||||
with pytest.raises(type(refusal)):
|
||||
await ps.ensure_spend_counter_initialized(
|
||||
counter_key="spend:end_user:customer-1", source_cache_key="customer-1"
|
||||
)
|
||||
|
||||
assert await redis.get("spend:end_user:customer-1") is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue