test(batches): check the batch usage counts in the e2e terminal-state test

tests/e2e/batches/batch_client.py

@@ -15,7 +15,7 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import Final, Literal
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 from e2e_metadata import step
 from proxy_client import ProxyClient
@@ -77,6 +77,16 @@ class BatchErrorList(BaseModel):
     data: list[BatchErrorItem] = []
 
 
+class BatchUsage(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    prompt_tokens: int | None = None
+    completion_tokens: int | None = None
+    total_tokens: int | None = None
+    input_tokens: int | None = None
+    output_tokens: int | None = None
+
+
 class BatchObject(BaseModel):
     id: str
     object: str | None = None
@@ -90,6 +100,7 @@ class BatchObject(BaseModel):
     completion_window: str | None = None
     created_at: int | None = None
     model: str | None = None
+    usage: BatchUsage | None = None
 
 
 class BatchList(BaseModel):

tests/e2e/batches/test_batches_e2e.py

@@ -1893,6 +1893,13 @@ class TestBatchTerminalState:
             f"listed-completed marker retrieved as {fetched.status!r}"
         )
         assert fetched.output_file_id, "completed batch has no output_file_id"
+        usage: Final = fetched.usage
+        assert usage is not None and usage.input_tokens and usage.output_tokens, (
+            f"completed batch has no batch token counts: {usage!r}"
+        )
+        assert (usage.prompt_tokens, usage.completion_tokens) == (usage.input_tokens, usage.output_tokens), (
+            f"chat token counts disagree with the batch token counts: {usage!r}"
+        )
 
         downloaded = client.proxy.transport.download(
             f"/v1/files/{fetched.output_file_id}/content",