fix(e2e): bound spend-log snapshots to a /spend/logs/v2 window

The rate-limited batch spend test snapshotted unattributed rows via the
unpaginated /spend/logs whole-table read, which grows with the environment
(58MB on stage) and OOMKilled the e2e runner at its 512Mi limit on every
scheduled run. Gateway.spend_logs_window pages /spend/logs/v2 over an
explicit date window instead, and SpendLogsParams now rejects a filterless
read so the whole-table call cannot come back
This commit is contained in:
mubashir1osmani 2026-07-14 12:56:35 -07:00
parent 2166608eb8
commit 8a5ffab1ca
4 changed files with 122 additions and 6 deletions

View file

@ -17,6 +17,7 @@ from __future__ import annotations
import json
import time
from datetime import datetime, timedelta, timezone
from typing import Callable
import pytest
@ -49,7 +50,7 @@ from e2e_http import (
unwrap,
)
from lifecycle import ResourceManager
from models import KeyGenerateBody, SpendLogRow, SpendLogsParams
from models import KeyGenerateBody, SpendLogRow
pytestmark = pytest.mark.e2e
@ -349,6 +350,10 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row(
the file-read path fires while the batch itself is not blocked.
``resources.key()`` cannot set limits, so the key is minted on the gateway
directly and its delete deferred.
Snapshots read /spend/logs/v2 over a bounded window around the test instead
of the unpaginated /spend/logs whole-table read, which grows with the
environment and OOMed the e2e runner on stage.
"""
user_id = f"e2e-batch-rl-{unique_marker()}"
key = client.gateway.generate_key(
@ -356,8 +361,13 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row(
)
resources.defer(lambda: client.gateway.delete_key(key))
window_start = datetime.now(timezone.utc) - timedelta(hours=1)
window_end = window_start + timedelta(hours=2)
before = frozenset(
row.request_id for row in unattributed_rows(client.gateway.spend_logs(SpendLogsParams()))
row.request_id
for row in unattributed_rows(
client.gateway.spend_logs_window(start=window_start, end=window_end)
)
)
file = unwrap(
@ -379,7 +389,9 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row(
new_orphans = [
row
for row in unattributed_rows(client.gateway.spend_logs(SpendLogsParams()))
for row in unattributed_rows(
client.gateway.spend_logs_window(start=window_start, end=window_end)
)
if row.request_id not in before
]
assert not new_orphans, (

View file

@ -12,6 +12,7 @@ import time
import warnings
from collections.abc import Callable
from dataclasses import dataclass
from datetime import datetime
from e2e_http import (
NoBody,
@ -50,6 +51,8 @@ from models import (
OcrResponse,
SpendLogRow,
SpendLogs,
SpendLogsPage,
SpendLogsPageParams,
SpendLogsParams,
)
from e2e_config import (
@ -255,6 +258,28 @@ class Gateway:
case _:
return []
def spend_logs_window(self, *, start: datetime, end: datetime) -> list[SpendLogRow]:
def fetch(page: int) -> SpendLogsPage:
return unwrap(
self.transport.get(
"/spend/logs/v2",
headers=self.transport.master,
params=SpendLogsPageParams(
start_date=start.strftime("%Y-%m-%d %H:%M:%S"),
end_date=end.strftime("%Y-%m-%d %H:%M:%S"),
page=page,
page_size=100,
),
response_type=SpendLogsPage,
)
)
first = fetch(1)
return [
*first.data,
*(row for page in range(2, first.total_pages + 1) for row in fetch(page).data),
]
def poll_logs_for_key(
self, key: str, *, min_rows: int = 1, predicate: RowsPredicate | None = None
) -> list[SpendLogRow]:

View file

@ -8,7 +8,7 @@ from __future__ import annotations
from typing import Literal
from pydantic import BaseModel, ConfigDict, RootModel
from pydantic import BaseModel, ConfigDict, RootModel, model_validator
# ---------- keys ----------
@ -255,6 +255,16 @@ class SpendLogsParams(BaseModel):
request_id: str | None = None
api_key: str | None = None
@model_validator(mode="after")
def require_filter(self) -> SpendLogsParams:
if self.request_id is None and self.api_key is None:
raise ValueError(
"unfiltered /spend/logs returns the entire spend table and OOMs the "
"runner on long-lived environments; filter by request_id or api_key, "
"or use Gateway.spend_logs_window for a bounded /spend/logs/v2 read"
)
return self
class SpendLogsPageParams(BaseModel):
"""Query for /spend/logs/v2, which requires an explicit date window and

View file

@ -1,17 +1,23 @@
"""Unit coverage for the Gateway model-management surface (create_model /
delete_model).
delete_model) and the bounded spend read-back (spend_logs_window).
The batches conftest and several llm_translation tests register deployments at
runtime through gateway.create_model; when that method went missing, every batch
test errored at fixture setup (AttributeError) before a single request reached
the proxy. This pins the surface with a typed fake Transport so a rename or
signature drift fails here instead of in a live stage run.
spend_logs_window exists because the unpaginated /spend/logs whole-table read
grew past the e2e runner's memory limit on stage and OOMKilled every run; these
tests pin its /spend/logs/v2 pagination and that SpendLogsParams can no longer
express the unfiltered read.
"""
from dataclasses import dataclass, field
from datetime import datetime, timezone
import pytest
from pydantic import BaseModel
from pydantic import BaseModel, ValidationError
from batches.batch_client import BatchClient
from e2e_gateway import Gateway
@ -30,6 +36,9 @@ from models import (
ModelNewBody,
ModelNewResponse,
ModelsListResponse,
SpendLogsPage,
SpendLogsPageParams,
SpendLogsParams,
)
@ -46,6 +55,8 @@ class _RecordingTransport:
servable_after_gets: int = 0
models_error: UnknownApiError | None = None
model_gets: int = 0
spend_total: int = 0
spend_gets: list[SpendLogsPageParams] = field(default_factory=list)
_created: list[str] = field(default_factory=list)
def post[R: BaseModel](
@ -91,6 +102,22 @@ class _RecordingTransport:
return Success(
data=response_type.model_validate({"data": [{"id": name} for name in visible]})
)
if path == "/spend/logs/v2" and response_type is SpendLogsPage:
assert isinstance(params, SpendLogsPageParams)
self.spend_gets.append(params)
offset = (params.page - 1) * params.page_size
count = min(params.page_size, max(self.spend_total - offset, 0))
return Success(
data=response_type.model_validate(
{
"data": [{"request_id": f"req-{offset + i}"} for i in range(count)],
"total": self.spend_total,
"page": params.page,
"page_size": params.page_size,
"total_pages": (self.spend_total + params.page_size - 1) // params.page_size,
}
)
)
raise AssertionError(f"unexpected get: {path}")
def delete[R: BaseModel](
@ -202,3 +229,45 @@ def test_gateway_delete_model_posts_the_model_id() -> None:
assert path == "/model/delete"
assert isinstance(body, ModelDeleteBody)
assert body.id == "registered-id"
WINDOW_START = datetime(2026, 7, 14, 12, 0, 0, tzinfo=timezone.utc)
WINDOW_END = datetime(2026, 7, 14, 14, 0, 0, tzinfo=timezone.utc)
def test_gateway_spend_logs_window_pages_through_every_row_in_the_window() -> None:
transport = _RecordingTransport(spend_total=250)
gateway = Gateway(transport=transport)
rows = gateway.spend_logs_window(start=WINDOW_START, end=WINDOW_END)
assert len(rows) == 250
assert len({row.request_id for row in rows}) == 250
assert [params.page for params in transport.spend_gets] == [1, 2, 3]
assert all(params.start_date == "2026-07-14 12:00:00" for params in transport.spend_gets)
assert all(params.end_date == "2026-07-14 14:00:00" for params in transport.spend_gets)
def test_gateway_spend_logs_window_stops_at_an_exact_page_boundary() -> None:
transport = _RecordingTransport(spend_total=200)
gateway = Gateway(transport=transport)
rows = gateway.spend_logs_window(start=WINDOW_START, end=WINDOW_END)
assert len(rows) == 200
assert [params.page for params in transport.spend_gets] == [1, 2]
def test_gateway_spend_logs_window_returns_empty_for_an_empty_window() -> None:
transport = _RecordingTransport(spend_total=0)
gateway = Gateway(transport=transport)
rows = gateway.spend_logs_window(start=WINDOW_START, end=WINDOW_END)
assert rows == []
assert [params.page for params in transport.spend_gets] == [1]
def test_spend_logs_params_rejects_the_unfiltered_whole_table_read() -> None:
with pytest.raises(ValidationError, match="spend_logs_window"):
SpendLogsParams()