From cec9797cb07a3e28fff7c25c4c4c7fe1aec949a3 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 23 Jul 2026 14:59:48 -0700 Subject: [PATCH] test(e2e/batches): cover GET /v1/batches pagination for openai + vertex --- tests/e2e/batches/COVERAGE.md | 15 +- tests/e2e/batches/batch_client.py | 17 +- tests/e2e/batches/capabilities.py | 1 + tests/e2e/batches/test_batches_e2e.py | 151 +++++++++++++++++- .../llm_nonconversational.yaml | 2 + 5 files changed, 182 insertions(+), 4 deletions(-) diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index ca48204962a..82fea6f4cbe 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -45,6 +45,19 @@ provider also fails create (the file id / model do not belong there), and the `provider_fallback` raw batch id is additionally checked against the provider's native shape (`raw_id_matches_provider`). +## List pagination + +`TestBatchListPagination` walks the managed-batch listing the way a client does: +it creates several batches under a key with a unique `user_id`, then pages +`GET /v1/batches` with `limit=1`, following each page's last id as the `after` +cursor until the proxy reports no more. The unique `user_id` scopes the +owner-filtered listing to exactly this test's batches, so the expected page order +is deterministic. It asserts every created id is returned exactly once, in +reverse-chronological (newest-first) order, and that the walk terminates rather +than looping. Covered for OpenAI and Vertex, both routed through the managed +listing path. This is the paginate-until-all-ids-found workflow, and it guards +against a cursor that repeats a page (loop) or drops batches (missing ids). + ## Key model restriction `test_batch_key_model_access_denied` mints a key restricted to one model @@ -71,7 +84,7 @@ File delete asserts `object=="file"` and `deleted==True`. | `batch_client.py` | typed file upload/download + batch create/retrieve/cancel/list/delete over the shared ProxyClient; runtime batch model registration via /model/new; denial helpers | | `capabilities.py` | the provider x scenario matrix + per-provider /model/new params + id-shape classifiers + per-provider raw-id assertion | | `conftest.py` | session-scoped batch deployment registration and teardown | -| `test_batches_e2e.py` | parametrized lifecycle with per-endpoint output assertions, file upload/delete outputs, key-model-access denial | +| `test_batches_e2e.py` | parametrized lifecycle with per-endpoint output assertions, file upload/delete outputs, key-model-access denial, list pagination (openai + vertex) | ## Out of scope (intentionally) diff --git a/tests/e2e/batches/batch_client.py b/tests/e2e/batches/batch_client.py index 5cc5d1dae3b..d38b5dca68d 100644 --- a/tests/e2e/batches/batch_client.py +++ b/tests/e2e/batches/batch_client.py @@ -59,6 +59,9 @@ class BatchObject(BaseModel): class BatchList(BaseModel): object: str | None = None data: list[BatchObject] = [] + has_more: bool | None = None + first_id: str | None = None + last_id: str | None = None class FileDeleteResponse(BaseModel): @@ -78,6 +81,11 @@ class ModelQuery(BaseModel): model: str | None = None +class BatchListQuery(BaseModel): + limit: int | None = None + after: str | None = None + + def is_model_access_denied(resp: StreamingResponse) -> bool: """True if the proxy rejected the call because the key may not access the model.""" return resp.status_code == 403 and "key_model_access_denied" in resp.body @@ -168,12 +176,17 @@ class BatchClient: ) def list_batches( - self, *, key: str, provider: str | None = None + self, + *, + key: str, + provider: str | None = None, + limit: int | None = None, + after: str | None = None, ) -> Result[BatchList]: return self.proxy.transport.get( _batches_path(provider), headers=self.proxy.transport.bearer(key), - params=NoBody(), + params=BatchListQuery(limit=limit, after=after), response_type=BatchList, ) diff --git a/tests/e2e/batches/capabilities.py b/tests/e2e/batches/capabilities.py index 3988fb5e7e1..c5c4249885d 100644 --- a/tests/e2e/batches/capabilities.py +++ b/tests/e2e/batches/capabilities.py @@ -144,6 +144,7 @@ def _model_for(provider_name: str) -> str: OPENAI_BATCH_MODEL = _model_for("openai") AZURE_BATCH_MODEL = _model_for("azure") +VERTEX_BATCH_MODEL = _model_for("vertex_ai") BEDROCK_SCENARIOS: tuple[Scenario, ...] = ("unified",) diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index b0c53becb6b..2377810b426 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -18,8 +18,9 @@ from __future__ import annotations import json import os import time +from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Callable +from typing import Callable, Iterator import pytest @@ -40,6 +41,7 @@ from capabilities import ( CAPABILITIES, FILE_ID_SHAPE, OPENAI_BATCH_MODEL, + VERTEX_BATCH_MODEL, Capability, batch_model_name, coverage_cells_for_lifecycle, @@ -852,3 +854,150 @@ class TestHostedVllmBatch: f"hosted_vllm batch has non-transitional status {batch.status!r}" ) assert_batch_object(batch) + + +BATCH_PAGE_COUNT = 5 +BATCH_PAGE_SIZE = 1 +# Newest-first, one batch per page: N data pages + one empty terminating page. +# The cap is generous headroom; only a non-advancing cursor (a loop) hits it. +BATCH_PAGE_FETCH_CAP = 2 * BATCH_PAGE_COUNT + 3 + + +@dataclass(frozen=True, slots=True) +class PaginationCase: + provider: str + deployment: str + jsonl_model: str + required_env: tuple[str, ...] + cell: str + + +PAGINATION_CASES: tuple[PaginationCase, ...] = ( + PaginationCase( + provider="openai", + deployment=OPENAI_BATCH_MODEL, + jsonl_model="gpt-4o-mini", + required_env=("OPENAI_API_KEY",), + cell="llm.batches.openai.list_pagination.nonstream.works", + ), + PaginationCase( + provider="vertex_ai", + deployment=VERTEX_BATCH_MODEL, + jsonl_model="gemini-2.5-flash", + required_env=("VERTEXAI_PROJECT", "VERTEXAI_CREDENTIALS", "GCS_BUCKET_NAME"), + cell="llm.batches.vertex.list_pagination.nonstream.works", + ), +) + + +def create_batch_resilient_unified( + client: BatchClient, file_id: str, key: str +) -> StreamingResponse: + last = client.create_batch(body=BatchCreateBody(input_file_id=file_id), key=key) + for attempt in range(BATCH_OP_RETRIES - 1): + if last.ok or not _transient_status(last.status_code): + return last + time.sleep(_backoff_seconds(attempt)) + last = client.create_batch(body=BatchCreateBody(input_file_id=file_id), key=key) + return last + + +def paginate_batch_ids( + client: BatchClient, key: str, page_size: int, cap: int +) -> Iterator[str]: + """Walk GET /v1/batches the way a client does: read a page, follow its last id + as the `after` cursor, repeat until the proxy reports no more. Yields the batch + ids in the order the proxy paged them. Raises if the cursor never advances past + `cap` pages, which is what a pagination loop looks like from the client side.""" + after: str | None = None + for _ in range(cap): + page = unwrap(client.list_batches(key=key, limit=page_size, after=after)) + if not page.data: + return + yield from (batch.id for batch in page.data) + if not page.has_more: + return + after = page.data[-1].id + raise AssertionError( + f"GET /v1/batches never terminated within {cap} pages (page_size={page_size}); " + "the `after` cursor is not advancing, so the client loops forever" + ) + + +class TestBatchListPagination: + """GET /v1/batches paginates the caller's batches newest-first without loss. + + A client that lists batches and walks the `after` cursor to find a known set + of batch ids must see every id exactly once, in reverse-chronological order, + and the walk must terminate. This mirrors a customer workflow that polls batch + status by paginating GET /v1/batches until every expected id is found. + + Isolation: batches are created under a key with a unique user_id, so the + owner-scoped managed-batch listing returns exactly this test's batches and the + expected page order is deterministic (the reverse of creation order). + """ + + @pytest.mark.parametrize( + "case", + [ + pytest.param(case, id=case.provider, marks=pytest.mark.covers(case.cell)) + for case in PAGINATION_CASES + ], + ) + def test_pagination_finds_every_batch_in_reverse_order( + self, + case: PaginationCase, + client: BatchClient, + resources: ResourceManager, + batch_deployments: None, + ) -> None: + require_env(*case.required_env) + + user_id = f"e2e-batch-page-{case.provider}-{unique_marker()}" + key = client.proxy.generate_key(KeyGenerateBody(models=[], user_id=user_id)) + resources.defer(lambda: client.proxy.delete_key(key)) + + file = unwrap( + client.upload_file( + content=render_jsonl(case.jsonl_model), + form=FileUploadForm(purpose="batch", target_model_names=case.deployment), + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + + def create_managed_batch() -> str: + created = create_batch_resilient_unified(client, file.id, key) + require_successful_call(created) + batch = BatchObject.model_validate_json(created.body) + resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + assert is_managed_id(batch.id), ( + f"{case.provider}: unified create must return a managed batch id so the " + f"managed listing indexes it, got {batch.id!r}" + ) + return batch.id + + created_ids = tuple(create_managed_batch() for _ in range(BATCH_PAGE_COUNT)) + assert len(set(created_ids)) == BATCH_PAGE_COUNT, ( + f"{case.provider}: create returned duplicate batch ids {created_ids}" + ) + expected = tuple(reversed(created_ids)) + + collected = tuple( + paginate_batch_ids(client, key, BATCH_PAGE_SIZE, BATCH_PAGE_FETCH_CAP) + ) + + missing = tuple(bid for bid in created_ids if bid not in collected) + assert not missing, ( + f"{case.provider}: paginating GET /v1/batches never returned " + f"{len(missing)} of the {BATCH_PAGE_COUNT} created batches: {missing}. " + f"pages yielded {collected}" + ) + assert len(collected) == len(set(collected)), ( + f"{case.provider}: pagination returned the same batch id on more than one " + f"page (a cursor loop): {collected}" + ) + assert collected == expected, ( + f"{case.provider}: GET /v1/batches must page newest-first. " + f"expected reverse-creation order {expected}, got {collected}" + ) diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index 63e6fde14a3..3c45b60049d 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -10,6 +10,7 @@ - {id: llm.batches.openai.retrieve.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Batch retrieve, id round-trip + status"} - {id: llm.batches.openai.cancel.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Batch cancel"} - {id: llm.batches.openai.list.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Batch list envelope"} +- {id: llm.batches.openai.list_pagination.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Managed-batch listing must page the caller's batches newest-first via the after cursor with no loss and no cursor loop"} - {id: llm.batches.openai.file_lifecycle.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "File upload/retrieve/delete for batch flow"} - {id: llm.batches.openai_encoded.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py", rationale: "Encoded scenario lifecycle"} - {id: llm.batches.openai_unified.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py", rationale: "Unified/managed-id scenario"} @@ -17,6 +18,7 @@ - {id: llm.batches.openai_provider_fallback.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py", rationale: "Provider-fallback raw-id scenario"} - {id: llm.batches.azure_openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:98", rationale: "Azure batches all scenarios"} - {id: llm.batches.vertex.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:98", rationale: "Vertex batches"} +- {id: llm.batches.vertex.list_pagination.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Vertex managed-batch listing must page newest-first via the after cursor with no loss and no cursor loop"} - {id: llm.batches.bedrock.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:98", rationale: "Bedrock batches (encoded/unified only)"} - {id: llm.batches.bedrock.assume_role.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: bedrock_converse, capability: assume_role, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Bedrock batch create under STS assume-role credentials"} - {id: llm.batches.hosted_vllm.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: hosted_vllm, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "hosted_vllm OpenAI-compatible batch create"}