scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 206

test(spend): run explained reads within real read transactions (#49996)

Browse source 
Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-11-04 05:11:23 +00:00 • committed by GitHub
parent 767613bd40
commit 0150896a50
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 14 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
tests/integration/spend/test_session_grouped_logs_pages.py
View file

@ -319,6 +319,18 @@ class _ExplainedReads:
self.database: Final = database
self.rows_scanned: tuple[int, ...] = ()
@asynccontextmanager
async def tx(self, *, timeout: timedelta) -> AsyncIterator["_ExplainedReads"]:
async with self.database.tx(timeout=timeout) as transaction:
reads: Final = _ExplainedReads(transaction)
try:
yield reads
finally:
self.rows_scanned = (*self.rows_scanned, *reads.rows_scanned)
async def execute_raw(self, query: str, *args: object) -> int:
return await self.database.execute_raw(query, *args)
async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]:
explained: Final = await self.database.query_raw(f"EXPLAIN (ANALYZE, FORMAT JSON) {query}", *args)
plan: Final = _Explained.model_validate(explained[0]["QUERY PLAN"][0]).plan
@ -384,14 +396,14 @@ async def test_a_page_cut_inside_a_tie_reads_on_until_the_tie_is_whole() -> None
@dataclass(frozen=True, slots=True)
class _ExplainedPrisma:
db: _ExplainedReads
read_db: _ExplainedReads
async def test_a_pages_representatives_are_found_by_session_not_by_scanning_the_window() -> None:
async with _spend_logs_copy(*_generated_sessions(rows_per_session=10)) as database:
reads: Final = _ExplainedReads(database)
representatives: Final = await _fetch_session_representatives(
prisma_client=_ExplainedPrisma(reads), # pyright: ignore[reportArgumentType] # only .db is used
prisma_client=_ExplainedPrisma(reads), # pyright: ignore[reportArgumentType] # only .read_db is used
where_clause=_COPY_WINDOW,
sql_params=_COPY_WINDOW_PARAMS,
next_param_index=len(_COPY_WINDOW_PARAMS) + 1,