Files
567-labs--instructor/tests/test_batch_processor_coverage.py
T
wehub-resource-sync 97e91a83f3
Ruff / Ruff (push) Waiting to run
Test / Core Tests (push) Waiting to run
Test / Offline Coverage Tests (Python 3.10) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.11) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.12) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.13) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.9) (push) Waiting to run
Test / Full Coverage (Python 3.11) (push) Waiting to run
Test / Core Provider Tests (OpenAI) (push) Blocked by required conditions
Test / Core Provider Tests (Anthropic) (push) Blocked by required conditions
Test / Core Provider Tests (Google) (push) Blocked by required conditions
Test / Core Provider Tests (Other) (push) Blocked by required conditions
Test / Anthropic Tests (push) Blocked by required conditions
Test / Gemini Tests (push) Blocked by required conditions
Test / Google GenAI Tests (push) Blocked by required conditions
Test / Vertex AI Tests (push) Blocked by required conditions
Test / OpenAI Tests (push) Blocked by required conditions
Test / Writer Tests (push) Blocked by required conditions
Test / Auto Client Tests (push) Blocked by required conditions
ty / type-check (push) Waiting to run
chore: import upstream snapshot with attribution
2026-07-13 13:36:38 +08:00

410 lines
13 KiB
Python

"""Focused tests for the public batch processor API and result parsing."""
from __future__ import annotations
import io
import json
from pathlib import Path
from typing import Any
import pytest
from pydantic import BaseModel
import instructor.batch.processor as processor_module
from instructor.batch.models import BatchError, BatchJobInfo, BatchSuccess
from instructor.batch.processor import BatchProcessor
pytestmark = pytest.mark.unit
class Person(BaseModel):
name: str
age: int
class RecordingProvider:
def __init__(self, results: str = "") -> None:
self.results = results
self.calls: list[tuple[str, Any]] = []
def submit_batch(
self,
file_path_or_buffer: str | io.BytesIO,
metadata: dict[str, Any] | None = None,
**kwargs: Any,
) -> str:
self.calls.append(("submit", (file_path_or_buffer, metadata, kwargs)))
return "batch-123"
def get_status(self, batch_id: str) -> dict[str, Any]:
self.calls.append(("status", batch_id))
return {"id": batch_id, "status": "completed"}
def retrieve_results(self, batch_id: str) -> str:
self.calls.append(("retrieve", batch_id))
return self.results
def download_results(self, batch_id: str, file_path: str) -> None:
self.calls.append(("download", (batch_id, file_path)))
Path(file_path).write_text(self.results)
def cancel_batch(self, batch_id: str) -> dict[str, Any]:
self.calls.append(("cancel", batch_id))
return {"id": batch_id, "status": "cancelled"}
def delete_batch(self, batch_id: str) -> dict[str, Any]:
self.calls.append(("delete", batch_id))
return {"id": batch_id, "deleted": True}
def list_batches(self, limit: int = 10) -> list[BatchJobInfo]:
self.calls.append(("list", limit))
return [
BatchJobInfo.from_openai(
{"id": "batch-123", "status": "completed", "metadata": {}}
)
]
@pytest.fixture
def provider(monkeypatch: pytest.MonkeyPatch) -> RecordingProvider:
provider = RecordingProvider()
monkeypatch.setattr(processor_module, "get_provider", lambda _: provider)
return provider
def openai_result(custom_id: str, content: str) -> str:
return json.dumps(
{
"custom_id": custom_id,
"response": {"body": {"choices": [{"message": {"content": content}}]}},
}
)
def test_init_splits_provider_and_model_and_rejects_invalid_model(
provider: RecordingProvider,
) -> None:
processor = BatchProcessor("openai/gpt-4.1-mini", Person)
assert processor.provider is provider
assert processor.provider_name == "openai"
assert processor.model_name == "gpt-4.1-mini"
assert processor.response_model is Person
with pytest.raises(
ValueError, match='Model string must be in format "provider/model-name"'
):
BatchProcessor("gpt-4.1-mini", Person)
def test_create_batch_file_replaces_existing_contents_and_serializes_requests(
provider: RecordingProvider, tmp_path: Path, capsys: pytest.CaptureFixture[str]
) -> None:
del provider
batch_file = tmp_path / "requests.jsonl"
batch_file.write_text("stale request\n")
processor = BatchProcessor("openai/gpt-4.1-mini", Person)
returned = processor.create_batch_from_messages(
[
[{"role": "user", "content": "Ada is 36"}],
[{"role": "user", "content": "Lin is 28"}],
],
str(batch_file),
max_tokens=321,
temperature=0.3,
)
lines = [json.loads(line) for line in batch_file.read_text().splitlines()]
assert returned == str(batch_file)
assert [line["custom_id"] for line in lines] == ["request-0", "request-1"]
assert [line["body"]["model"] for line in lines] == [
"gpt-4.1-mini",
"gpt-4.1-mini",
]
assert lines[0]["body"]["max_tokens"] == 321
assert lines[0]["body"]["temperature"] == 0.3
assert "Created batch file" in capsys.readouterr().out
def test_create_batch_file_creates_a_new_output_file(
provider: RecordingProvider, tmp_path: Path
) -> None:
del provider
batch_file = tmp_path / "new-requests.jsonl"
processor = BatchProcessor("openai/gpt-4.1-mini", Person)
returned = processor.create_batch_from_messages(
[[{"role": "user", "content": "Ada is 36"}]], str(batch_file)
)
assert returned == str(batch_file)
request = json.loads(batch_file.read_text())
assert request["custom_id"] == "request-0"
assert request["body"]["messages"] == [{"role": "user", "content": "Ada is 36"}]
def test_create_batch_buffer_is_readable_from_start(
provider: RecordingProvider, capsys: pytest.CaptureFixture[str]
) -> None:
del provider
processor = BatchProcessor("anthropic/claude-sonnet", Person)
buffer = processor.create_batch_from_messages(
[[{"role": "user", "content": "Ada is 36"}]], max_tokens=90, temperature=0.2
)
assert isinstance(buffer, io.BytesIO)
assert buffer.tell() == 0
request = json.loads(buffer.read().decode())
assert request["custom_id"] == "request-0"
assert request["params"]["model"] == "claude-sonnet"
assert request["params"]["max_tokens"] == 90
assert request["params"]["temperature"] == 0.2
assert "Created batch buffer with 1 requests" in capsys.readouterr().out
def test_provider_operations_forward_arguments_and_parse_downloaded_results(
provider: RecordingProvider, tmp_path: Path
) -> None:
provider.results = openai_result("request-0", '{"name": "Ada", "age": 36}')
processor = BatchProcessor("openai/gpt-4.1-mini", Person)
request_buffer = io.BytesIO(b"request")
assert (
processor.submit_batch(request_buffer, completion_window="24h") == "batch-123"
)
assert provider.calls[-1] == (
"submit",
(
request_buffer,
{"description": "Instructor batch job"},
{"completion_window": "24h"},
),
)
assert processor.submit_batch("requests.jsonl", metadata={"team": "search"}) == (
"batch-123"
)
assert provider.calls[-1] == (
"submit",
("requests.jsonl", {"team": "search"}, {}),
)
assert processor.get_batch_status("batch-123") == {
"id": "batch-123",
"status": "completed",
}
jobs = processor.list_batches(limit=2)
assert len(jobs) == 1
assert jobs[0].id == "batch-123"
assert provider.calls[-1] == ("list", 2)
results_file = tmp_path / "results.jsonl"
results = processor.get_results("batch-123", str(results_file))
assert isinstance(results[0], BatchSuccess)
assert results[0].result == Person(name="Ada", age=36)
assert results_file.read_text() == provider.results
assert ("retrieve", "batch-123") in provider.calls
assert ("download", ("batch-123", str(results_file))) in provider.calls
calls_before = len(provider.calls)
in_memory_results = processor.get_results("batch-123")
assert isinstance(in_memory_results[0], BatchSuccess)
assert provider.calls[calls_before:] == [("retrieve", "batch-123")]
assert processor.cancel_batch("batch-123") == {
"id": "batch-123",
"status": "cancelled",
}
assert processor.delete_batch("batch-123") == {
"id": "batch-123",
"deleted": True,
}
assert provider.calls == [
(
"submit",
(
request_buffer,
{"description": "Instructor batch job"},
{"completion_window": "24h"},
),
),
("submit", ("requests.jsonl", {"team": "search"}, {})),
("status", "batch-123"),
("list", 2),
("retrieve", "batch-123"),
("download", ("batch-123", str(results_file))),
("retrieve", "batch-123"),
("cancel", "batch-123"),
("delete", "batch-123"),
]
def test_openai_results_distinguish_success_validation_extraction_and_json_errors(
provider: RecordingProvider,
) -> None:
del provider
processor = BatchProcessor("openai/gpt-4.1-mini", Person)
content = "\n".join(
[
openai_result("ok", '{"name": "Ada", "age": 36}'),
" ",
openai_result("invalid-model", '{"name": "Ada", "age": "unknown"}'),
json.dumps({"custom_id": "missing-response"}),
"not-json",
]
)
results = processor.parse_results(content)
assert len(results) == 4
assert isinstance(results[0], BatchSuccess)
assert results[0].custom_id == "ok"
assert results[0].result == Person(name="Ada", age=36)
assert isinstance(results[1], BatchError)
assert results[1].custom_id == "invalid-model"
assert results[1].error_type == "parsing_error"
assert "Failed to parse into Person" in results[1].error_message
assert results[1].raw_data == {"name": "Ada", "age": "unknown"}
assert isinstance(results[2], BatchError)
assert results[2].custom_id == "missing-response"
assert results[2].error_type == "extraction_error"
assert results[2].error_message == "Unknown error"
assert isinstance(results[3], BatchError)
assert results[3].custom_id == "unknown"
assert results[3].error_type == "json_parse_error"
assert results[3].raw_data == {"raw_line": "not-json"}
def test_anthropic_results_support_tool_use_and_text_fallback(
provider: RecordingProvider,
) -> None:
del provider
processor = BatchProcessor("anthropic/claude-sonnet", Person)
tool_result = {
"custom_id": "tool",
"result": {
"type": "succeeded",
"message": {
"content": [
{"type": "text", "text": "not json"},
{"type": "tool_use", "input": {"name": "Ada", "age": 36}},
]
},
},
}
text_result = {
"custom_id": "text",
"result": {
"type": "succeeded",
"message": {
"content": [
{"type": "text", "text": "not json"},
{"type": "text", "text": '{"name": "Lin", "age": 28}'},
]
},
},
}
results = processor.parse_results(
"\n".join([json.dumps(tool_result), json.dumps(text_result)])
)
assert len(results) == 2
assert isinstance(results[0], BatchSuccess)
assert isinstance(results[1], BatchSuccess)
assert results[0].custom_id == "tool"
assert results[0].result == Person(name="Ada", age=36)
assert results[1].custom_id == "text"
assert results[1].result == Person(name="Lin", age=28)
@pytest.mark.parametrize(
("result", "expected_type", "expected_message"),
[
(
{
"type": "error",
"error": {
"error": {"type": "rate_limit_error", "message": "slow down"}
},
},
"rate_limit_error",
"slow down",
),
(
{"type": "error", "error": "service unavailable"},
"anthropic_error",
"service unavailable",
),
(
{"type": "succeeded", "message": {"content": []}},
"extraction_error",
"Unknown error",
),
],
)
def test_anthropic_error_and_empty_results_keep_provider_details(
provider: RecordingProvider,
result: dict[str, Any],
expected_type: str,
expected_message: str,
) -> None:
del provider
processor = BatchProcessor("anthropic/claude-sonnet", Person)
parsed = processor.parse_results(
json.dumps({"custom_id": "request-7", "result": result})
)
assert len(parsed) == 1
assert isinstance(parsed[0], BatchError)
assert parsed[0].custom_id == "request-7"
assert parsed[0].error_type == expected_type
assert parsed[0].error_message == expected_message
def test_extract_returns_none_for_missing_malformed_or_unknown_provider_responses(
provider: RecordingProvider,
) -> None:
del provider
anthropic = BatchProcessor("anthropic/claude-sonnet", Person)
openai = BatchProcessor("openai/gpt-4.1-mini", Person)
unknown = BatchProcessor("local/model", Person)
assert anthropic._extract_from_response({"custom_id": "no-result"}) is None
assert (
anthropic._extract_from_response(
{"result": {"type": "succeeded", "message": {"content": "not-a-list"}}}
)
is None
)
assert anthropic._extract_from_response({"result": {"type": "succeeded"}}) is None
assert (
anthropic._extract_from_response(
{
"result": {
"type": "succeeded",
"message": {"content": [{"type": "image", "source": {}}]},
}
}
)
is None
)
assert anthropic._extract_from_response(
{
"result": {
"type": "succeeded",
"message": {
"content": [
{"type": "image", "source": {}},
{"type": "text", "text": '{"name":"Ada","age":36}'},
]
},
}
}
) == {"name": "Ada", "age": 36}
assert (
openai._extract_from_response({"response": {"body": {"choices": []}}}) is None
)
assert unknown._extract_from_response({"anything": "goes"}) is None