diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 7d54f05656e..85d9315b8c6 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -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, ( diff --git a/tests/e2e/e2e_gateway.py b/tests/e2e/e2e_gateway.py index 05f83ecc085..d40b96d60fa 100644 --- a/tests/e2e/e2e_gateway.py +++ b/tests/e2e/e2e_gateway.py @@ -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]: diff --git a/tests/e2e/models.py b/tests/e2e/models.py index e32f2709181..c8fc6c1ad4d 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -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 diff --git a/tests/e2e/test_e2e_gateway.py b/tests/e2e/test_e2e_gateway.py index 9a9aa2fd2cc..70e846ad8d5 100644 --- a/tests/e2e/test_e2e_gateway.py +++ b/tests/e2e/test_e2e_gateway.py @@ -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()