spb/doc-api

1"""Gemini Batch API. List and bogus-name probes always run; creating a batch is gated by RUN_BATCH_TESTS=true2(marker run_batch_tests) and xfails on FAILED_PRECONDITION (Batch API is not available on the free tier — observed 2026-09-18)."""3from __future__ import annotations45import time6import pytest78TERMINAL = {"BATCH_STATE_SUCCEEDED", "BATCH_STATE_FAILED", "BATCH_STATE_CANCELLED", "BATCH_STATE_EXPIRED"}91011def test_list_batches(gemini):12    st, body, _ = gemini("GET", "/v1beta/batches?pageSize=5", note="test_batch list")13    assert st == 200 and isinstance(body, dict)14    for op in body.get("operations", []):15        assert op["name"].startswith("batches/") and "done" in op161718def test_get_bogus_batch(gemini):19    st, body, _ = gemini("GET", "/v1beta/batches/bogus-batch", note="test_batch bogus")20    assert st in (400, 403, 404) and "error" in body212223@pytest.mark.run_batch_tests24def test_inline_batch_lifecycle(gemini, models):25    req = {"request": {"contents": [{"parts": [{"text": "Reply with OK."}]}], "generationConfig": {"maxOutputTokens": 8}}, "metadata": {"key": "r1"}}26    body = {"batch": {"display_name": "atlas-test", "input_config": {"requests": {"requests": [req]}}}}27    st, op, _ = gemini("POST", f"/v1beta/models/{models['gemini']}:batchGenerateContent", body, est_cost_usd=0.00005, note="test_batch create")28    if st == 400 and op["error"]["status"] == "FAILED_PRECONDITION":29        pytest.xfail("Batch API not enabled for this key (no free tier): " + op["error"]["message"])30    assert st == 200 and op["name"].startswith("batches/"), op31    assert op["metadata"]["state"] == "BATCH_STATE_PENDING"32    name = op["name"]33    try:34        t0 = time.time()35        state = op["metadata"]["state"]36        while not op.get("done") and time.time() - t0 < 300:37            time.sleep(20)38            st, op, _ = gemini("GET", f"/v1beta/{name}", note="test_batch poll")39            assert st == 20040            state = op["metadata"]["state"]41        if not op.get("done"):42            st, _, _ = gemini("POST", f"/v1beta/{name}:cancel", note="test_batch cancel")43            assert st == 20044        else:45            assert state in TERMINAL46            if state == "BATCH_STATE_SUCCEEDED":47                resp = op["response"]48                assert "inlinedResponses" in resp or "responsesFile" in resp49                if "inlinedResponses" in resp:50                    r0 = resp["inlinedResponses"]["inlinedResponses"][0]51                    assert r0["metadata"]["key"] == "r1" and ("response" in r0 or "error" in r0)52    finally:53        st, _, _ = gemini("DELETE", f"/v1beta/{name}", note="test_batch delete")54        assert st == 20055