From 385b4ff4640d7e4214458c6ac59646151b79cae5 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sun, 21 Jun 2026 00:36:13 +0000
Subject: [PATCH] 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.
---
 tests/e2e/budgets/test_spend_counter_reseed_e2e.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/tests/e2e/budgets/test_spend_counter_reseed_e2e.py b/tests/e2e/budgets/test_spend_counter_reseed_e2e.py
index 97bf0bae422..ab4a713543c 100644
--- a/tests/e2e/budgets/test_spend_counter_reseed_e2e.py
+++ b/tests/e2e/budgets/test_spend_counter_reseed_e2e.py
@@ -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: