mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
test(e2e): type the redis spend-counter probe for strict basedpyright
The new cold-counter reseed test drove its redis client untyped, so the strict tests/pyrightconfig.json (reportUnknown*, reportAny) flagged ten errors once the file landed: scan_iter/get came back unknown and the pool.map lambda had an untyped parameter. Annotate the client as redis.Redis[str] via a TYPE_CHECKING import (the runtime import stays lazy so the suite still skips, not errors, when redis is absent), which resolves scan_iter to Iterator[str] and get to str | None, and replace the lambda with a typed inner function mirroring _burst. basedpyright --project tests is back to zero errors.
This commit is contained in:
parent
74d5d1c3a7
commit
385b4ff464
1 changed files with 10 additions and 3 deletions
|
|
@ -20,6 +20,7 @@ import os
|
|||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from threading import Barrier
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -28,6 +29,9 @@ from e2e_config import unique_marker
|
|||
from e2e_http import StreamingResponse
|
||||
from lifecycle import ResourceManager
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import redis
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
MODEL = "claude-haiku-4-5"
|
||||
|
|
@ -38,7 +42,7 @@ BURST = 6
|
|||
COLD_WAIT_SECONDS = 80
|
||||
|
||||
|
||||
def _redis():
|
||||
def _redis() -> "redis.Redis[str]":
|
||||
import redis
|
||||
|
||||
return redis.Redis(
|
||||
|
|
@ -50,7 +54,7 @@ def _redis():
|
|||
)
|
||||
|
||||
|
||||
def _spend_counter(rds, key: str) -> float | None:
|
||||
def _spend_counter(rds: "redis.Redis[str]", key: str) -> float | None:
|
||||
"""The shared spend counter for `key`, or None if it is cold. The counter key is
|
||||
``{cache namespace}:spend:key:{sha256(key)}``; matched by suffix so the configured
|
||||
namespace need not be hard-coded."""
|
||||
|
|
@ -67,8 +71,11 @@ def _chat(client: BudgetClient, key: str) -> StreamingResponse:
|
|||
|
||||
|
||||
def _accumulate(client: BudgetClient, key: str, count: int) -> None:
|
||||
def one(_: int) -> StreamingResponse:
|
||||
return _chat(client, key)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
list(pool.map(lambda _: _chat(client, key), range(count)))
|
||||
list(pool.map(one, range(count)))
|
||||
|
||||
|
||||
def _burst(client: BudgetClient, key: str, count: int) -> None:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue