From ad57c5d1c2c144f3cb751c7e7058d6b8096ba49a Mon Sep 17 00:00:00 2001
From: mubashir1osmani <mubashir.osmani777@gmail.com>
Date: Thu, 16 Jul 2026 16:48:46 -0700
Subject: [PATCH] test(e2e): mark batches suite with coverage registry cells

Wire @pytest.mark.covers so non_core_llms counts the existing batches
lifecycle, key-access-denied, and file upload/delete tests against the
registry cells in llm_nonconversational.yaml
---
 tests/e2e/batches/capabilities.py     | 40 +++++++++++++++++++++++++++
 tests/e2e/batches/test_batches_e2e.py | 18 +++++++++++-
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/tests/e2e/batches/capabilities.py b/tests/e2e/batches/capabilities.py
index 3eb0e0328be..59097b70ef1 100644
--- a/tests/e2e/batches/capabilities.py
+++ b/tests/e2e/batches/capabilities.py
@@ -180,3 +180,43 @@ def matches_id_shape(shape: IdShape, id_str: str) -> bool:
     if shape == "model_encoded":
         return is_model_encoded_id(id_str)
     return not is_managed_id(id_str) and not is_model_encoded_id(id_str)
+
+
+def coverage_cells_for_lifecycle(cap: Capability) -> tuple[str, ...]:
+    """Registry cell ids that the parametrized lifecycle test covers for one capability.
+
+    OpenAI has per-scenario cells plus granular create/retrieve/cancel/list/file
+    cells. Other providers have one basic cell each. File-upload cells for the
+    batch-backing path are included when the lifecycle uploads for that provider.
+    """
+    match cap.provider:
+        case "openai":
+            cells = (
+                f"llm.batches.openai_{cap.scenario}.basic.nonstream.works",
+                "llm.batches.openai.create.nonstream.works",
+                "llm.batches.openai.retrieve.nonstream.works",
+                "llm.batches.openai.file_lifecycle.nonstream.works",
+                "llm.files.openai.upload.nonstream.works",
+            )
+            if cap.can_cancel:
+                cells = (*cells, "llm.batches.openai.cancel.nonstream.works")
+            if cap.can_list:
+                cells = (*cells, "llm.batches.openai.list.nonstream.works")
+            return cells
+        case "azure":
+            return (
+                "llm.batches.azure_openai.basic.nonstream.works",
+                "llm.files.azure_openai.upload.nonstream.works",
+            )
+        case "vertex_ai":
+            return (
+                "llm.batches.vertex.basic.nonstream.works",
+                "llm.files.vertex.upload.nonstream.works",
+            )
+        case "bedrock":
+            return (
+                "llm.batches.bedrock.basic.nonstream.works",
+                "llm.files.bedrock.upload.nonstream.works",
+            )
+        case _:
+            return ()
diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py
index 85d9315b8c6..2ee7eb36a41 100644
--- a/tests/e2e/batches/test_batches_e2e.py
+++ b/tests/e2e/batches/test_batches_e2e.py
@@ -37,6 +37,7 @@ from capabilities import (
     CAPABILITIES,
     FILE_ID_SHAPE,
     Capability,
+    coverage_cells_for_lifecycle,
     matches_id_shape,
     raw_id_matches_provider,
 )
@@ -168,7 +169,17 @@ def assert_batch_object(batch: BatchObject) -> None:
     ), "batch.created_at missing"
 
 
-@pytest.mark.parametrize("cap", CAPABILITIES, ids=[c.id for c in CAPABILITIES])
+@pytest.mark.parametrize(
+    "cap",
+    [
+        pytest.param(
+            cap,
+            id=cap.id,
+            marks=pytest.mark.covers(*coverage_cells_for_lifecycle(cap)),
+        )
+        for cap in CAPABILITIES
+    ],
+)
 def test_batch_lifecycle(
     cap: Capability,
     client: BatchClient,
@@ -266,6 +277,7 @@ def test_batch_lifecycle(
         assert match.object == "batch"
 
 
+@pytest.mark.covers("llm.batches.openai.key_model_access_denied.nonstream.works")
 def test_batch_key_model_access_denied(
     client: BatchClient, resources: ResourceManager, batch_deployments: None
 ) -> None:
@@ -301,6 +313,10 @@ def test_batch_key_model_access_denied(
     ), f"restricted key created a batch for a disallowed model (status {denied_create.status_code})"
 
 
+@pytest.mark.covers(
+    "llm.files.openai.upload.nonstream.works",
+    "llm.files.openai.delete.nonstream.works",
+)
 def test_file_upload_and_delete_outputs(
     client: BatchClient, resources: ResourceManager, batch_deployments: None
 ) -> None: