fix test

tests/test_litellm/_fixture_recorder.py

@@ -5,6 +5,7 @@ import hashlib
 import queue
 import threading
 from collections.abc import Callable, Generator
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
 from dataclasses import dataclass
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
@@ -51,6 +52,7 @@ class RecorderResult:
 
 @dataclass(frozen=True, slots=True)
 class GeneratorArgs:
+    concurrency: int
     examples: int
     fixture_dir: Path | None
     model: str
@@ -66,7 +68,7 @@ def _excluded_headers(headers: tuple[tuple[str, str], ...]) -> frozenset[str]:
 
 def _end_to_end_headers(headers: httpx.Headers) -> tuple[HttpHeader, ...]:
     decoded: Final = tuple((name.decode("ascii"), value.decode("latin-1")) for name, value in headers.raw)
-    excluded: Final = _excluded_headers(decoded) | {"content-length"}
+    excluded: Final = _excluded_headers(decoded) | {"content-encoding", "content-length"}
     return tuple(HttpHeader(name=name, value=value) for name, value in decoded if name.lower() not in excluded)
 
 
@@ -110,7 +112,7 @@ class _RecordingHandler(BaseHTTPRequestHandler):
                 content=request_body,
                 timeout=120,
             ) as upstream:
-                response_body: Final = b"".join(upstream.iter_raw())
+                response_body: Final = b"".join(upstream.iter_bytes())
                 recorded_response: Final = RecordedHttpResponse.from_bytes(
                     status_code=upstream.status_code,
                     headers=_end_to_end_headers(upstream.headers),
@@ -172,18 +174,39 @@ def record_case(
         sdk_call(api_base=recorder.url, api_key=spec.api_key, **case_input.as_sdk_kwargs())
         upstream_response: Final = recorder.take_response()
 
-    case: Final = OcrParityCase(input=case_input, upstream_response=upstream_response)
+    case: Final = OcrParityCase(litellm_input=case_input, provider_response=upstream_response)
     cache.put(cache_key, cast(dict[str, object], case.model_dump(mode="json", exclude_unset=True)))
     return RecorderResult(case=case, cache_hit=False)
 
 
+def record_cases(
+    spec: ProviderSpec,
+    root: Path,
+    case_inputs: tuple[MistralOcrParityInput, ...],
+    sdk_call: Callable[..., object],
+    max_concurrency: int,
+) -> tuple[RecorderResult, ...]:
+    if max_concurrency < 1:
+        raise ValueError("max_concurrency must be at least 1")
+    unique_inputs: Final = tuple(
+        {canonical_json(fixture_cache_key(case_input)): case_input for case_input in case_inputs}.values()
+    )
+    with ThreadPoolExecutor(max_workers=max_concurrency) as executor:
+        futures: Final = tuple(
+            executor.submit(record_case, spec, root, case_input, sdk_call) for case_input in unique_inputs
+        )
+        return tuple(future.result() for future in futures)
+
+
 def parse_generator_args() -> GeneratorArgs:
     parser: Final = argparse.ArgumentParser()
+    parser.add_argument("--concurrency", type=int, default=4)
     parser.add_argument("--examples", type=int, default=4)
     parser.add_argument("--fixture-dir", type=Path)
     parser.add_argument("--model", default="mistral/mistral-ocr-latest")
     namespace: Final = parser.parse_args()
     return GeneratorArgs(
+        concurrency=cast(int, namespace.concurrency),
         examples=cast(int, namespace.examples),
         fixture_dir=cast(Path | None, namespace.fixture_dir),
         model=cast(str, namespace.model),
@@ -202,13 +225,13 @@ def recorded_fixtures(directory: Path) -> tuple[OcrParityCase, ...]:
             fixtures.append(OcrParityCase.model_validate(raw_fixture))
         except ValidationError as error:
             raise ValueError(
-                f"invalid OCR parity fixture {path}: expected exactly `input` and `upstream_response` "
+                f"invalid OCR parity fixture {path}: expected exactly `litellm_input` and `provider_response` "
                 f"({len(error.errors())} validation errors)"
             ) from error
     return tuple(fixtures)
 
 
 def fixture_id(fixture: OcrParityCase) -> str:
-    input_json: Final = canonical_json(fixture.input.canonical_input())
+    input_json: Final = canonical_json(fixture.litellm_input.canonical_input())
     digest: Final = hashlib.sha256(input_json.encode("utf-8")).hexdigest()[:8]
-    return f"mistral-{fixture.input.model.rsplit('/', 1)[-1]}-{digest}"
+    return f"mistral-{fixture.litellm_input.model.rsplit('/', 1)[-1]}-{digest}"

tests/test_litellm/ocr/conftest.py

@@ -30,7 +30,7 @@ def pytest_generate_tests(metafunc: pytest.Metafunc) -> None:
     except (ValidationError, ValueError) as error:
         raise pytest.UsageError(
             f"Invalid OCR parity fixture bundle at {directory}. "
-            "Each fixture must contain exactly `input` and `upstream_response`. "
+            "Each fixture must contain exactly `litellm_input` and `provider_response`. "
             "Record fresh fixtures in an empty directory with: "
             f"`uv run python tests/test_litellm/ocr/generate_fixtures.py --fixture-dir {directory}`. "
             f"Validation details: {error}"

tests/test_litellm/ocr/fixture_models.py

@@ -87,5 +87,5 @@ class RecordedHttpResponse(_FixtureModel):
 
 
 class OcrParityCase(_FixtureModel):
-    input: MistralOcrParityInput
-    upstream_response: RecordedHttpResponse
+    litellm_input: MistralOcrParityInput
+    provider_response: RecordedHttpResponse

tests/test_litellm/ocr/generate_fixtures.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import logging
 import os
+import queue
 from collections.abc import Callable
 from pathlib import Path
 from typing import Final, cast
@@ -18,13 +19,10 @@ from tests.test_litellm._fixture_recorder import (
     ProviderSpec,
     fixture_directory,
     parse_generator_args,
-    record_case,
+    record_cases,
 )
 from tests.test_litellm.ocr.fixture_models import (
-    AnnotationFormat,
-    DocumentUrlDocument,
     ImageUrlDocument,
-    JsonSchemaDefinition,
     MistralOcrParityInput,
 )
 
@@ -32,16 +30,6 @@ FIXTURE_DIR_ENV: Final = "LITELLM_OCR_FIXTURE_DIR"
 LOGGER: Final = logging.getLogger(__name__)
 _TEXT: Final = st.from_regex(r"[A-Za-z0-9 ]{1,24}", fullmatch=True)
 _VALUE_TEXT: Final = st.text(alphabet="abcdefghijklmnopqrstuvwxyz0123456789 -_", min_size=1, max_size=32)
-_ANNOTATION_FORMAT: Final[SearchStrategy[AnnotationFormat]] = st.builds(
-    AnnotationFormat,
-    type=st.just("json_schema"),
-    json_schema=st.builds(
-        JsonSchemaDefinition,
-        name=_VALUE_TEXT,
-        schema_value=st.just({"type": "object", "properties": {}, "additionalProperties": False}),
-        strict=st.one_of(st.none(), st.booleans()),
-    ),
-)
 
 
 def _image_document(text: str, font_size: int) -> ImageUrlDocument:
@@ -50,36 +38,18 @@ def _image_document(text: str, font_size: int) -> ImageUrlDocument:
 
 
 def _input_strategy(model: str) -> SearchStrategy[MistralOcrParityInput]:
-    document_strategy: Final = st.one_of(
-        st.builds(_image_document, _TEXT, st.integers(min_value=12, max_value=36)),
-        st.just(
-            DocumentUrlDocument(
-                type="document_url",
-                document_url="https://arxiv.org/pdf/2201.04234",
-            )
-        ),
-    )
+    document_strategy: Final = st.builds(_image_document, _TEXT, st.integers(min_value=12, max_value=36))
     input_values: Final = st.fixed_dictionaries(
         {"model": st.just(model), "document": document_strategy},
         optional={
-            "pages": st.one_of(
-                st.none(),
-                st.lists(st.integers(min_value=0, max_value=20), min_size=1, max_size=5, unique=True),
-            ),
-            "include_image_base64": st.one_of(st.none(), st.booleans()),
-            "image_limit": st.one_of(st.none(), st.integers(min_value=1, max_value=100)),
-            "image_min_size": st.one_of(st.none(), st.integers(min_value=0, max_value=10_000)),
-            "bbox_annotation_format": st.one_of(st.none(), _ANNOTATION_FORMAT),
-            "document_annotation_format": st.one_of(st.none(), _ANNOTATION_FORMAT),
-            "document_annotation_prompt": st.one_of(st.none(), _VALUE_TEXT),
-            "extract_header": st.one_of(st.none(), st.booleans()),
-            "extract_footer": st.one_of(st.none(), st.booleans()),
-            "table_format": st.one_of(st.none(), st.sampled_from(("markdown", "html"))),
-            "confidence_scores_granularity": st.one_of(
-                st.none(), st.sampled_from(("word", "page", "block"))
-            ),
-            "include_blocks": st.one_of(st.none(), st.booleans()),
-            "id": st.one_of(st.none(), _VALUE_TEXT),
+            "include_image_base64": st.booleans(),
+            "image_limit": st.integers(min_value=1, max_value=100),
+            "image_min_size": st.integers(min_value=0, max_value=10_000),
+            "extract_header": st.booleans(),
+            "extract_footer": st.booleans(),
+            "table_format": st.sampled_from(("markdown", "html")),
+            "include_blocks": st.booleans(),
+            "id": _VALUE_TEXT,
         },
     )
     return input_values.map(MistralOcrParityInput.model_validate)
@@ -89,15 +59,22 @@ def _generate_examples(
     spec: ProviderSpec,
     root: Path,
     examples: int,
+    concurrency: int,
     sdk_call: Callable[..., object],
 ) -> None:
+    generated: Final[queue.SimpleQueue[MistralOcrParityInput | None]] = queue.SimpleQueue()
+
     @settings(max_examples=examples, deadline=None, derandomize=True)
     @given(case_input=_input_strategy(spec.model))
     def generate_case(case_input: MistralOcrParityInput) -> None:
-        result: Final = record_case(spec, root, case_input, sdk_call)
-        LOGGER.info("%s %s", "cached" if result.cache_hit else "recorded", result.case.input.model)
+        generated.put(case_input)
 
     generate_case()
+    generated.put(None)
+    case_inputs: Final = tuple(iter(generated.get, None))
+    results: Final = record_cases(spec, root, case_inputs, sdk_call, concurrency)
+    for result in results:
+        LOGGER.info("%s %s", "cached" if result.cache_hit else "recorded", result.case.litellm_input.model)
 
 
 def _mistral_upstream_base() -> str:
@@ -119,7 +96,7 @@ def main() -> None:
     )
     spec: Final = ProviderSpec(model=args.model, upstream_base=_mistral_upstream_base(), api_key=api_key)
     use_litellm_rust(False, ocr=None, aocr=None)
-    _generate_examples(spec, root, args.examples, cast(Callable[..., object], litellm.ocr))
+    _generate_examples(spec, root, args.examples, args.concurrency, cast(Callable[..., object], litellm.ocr))
 
 
 if __name__ == "__main__":

tests/test_litellm/ocr/test_sdk_parity.py

@@ -51,7 +51,7 @@ def _execute_sdk_case(sdk_input: MistralOcrParityInput, route: SDKRoute, mock_ur
 def test_recorded_ocr_sdk_parity(ocr_fixture: OcrParityCase, route: SDKRoute, tmp_path: Path) -> None:
     case_file: Final = tmp_path / f"{route.value}-ocr-parity-case.json"
     case_file.write_text(ocr_fixture.model_dump_json(indent=2, exclude_unset=True), encoding="utf-8")
-    response: Final = ocr_fixture.upstream_response
+    response: Final = ocr_fixture.provider_response
     response_body: Final = response.body_bytes()
     response_headers: Final = tuple((header.name, header.value) for header in response.headers)
     runner: Final = PythonScriptRunner(
@@ -90,7 +90,7 @@ def _child_main() -> None:
     mock_url: Final = sys.argv[3]
     report_file: Final = Path(sys.argv[4])
     case: Final = OcrParityCase.model_validate_json(case_file.read_text(encoding="utf-8"))
-    report: Final = _execute_sdk_case(case.input, route, mock_url)
+    report: Final = _execute_sdk_case(case.litellm_input, route, mock_url)
     report_file.write_text(report.model_dump_json(indent=2), encoding="utf-8")
 
 

tests/test_litellm/test__fixture_recorder.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+import threading
+from collections.abc import Generator
+from contextlib import contextmanager
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+from typing import Final, cast
+
+import httpx
+
+from tests.test_litellm._fixture_recorder import ProviderSpec, record_cases
+from tests.test_litellm.ocr.fixture_models import ImageUrlDocument, MistralOcrParityInput
+
+
+class _ControlledUpstream(ThreadingHTTPServer):
+    daemon_threads = True
+
+    def __init__(self) -> None:
+        super().__init__(("127.0.0.1", 0), _ControlledUpstreamHandler)
+        self.lock: Final = threading.Lock()
+        self.two_requests_started: Final = threading.Event()
+        self.active_requests: int = 0
+        self.max_active_requests: int = 0
+        self.request_count: int = 0
+
+    @property
+    def url(self) -> str:
+        return f"http://127.0.0.1:{self.server_address[1]}"
+
+    def start_request(self) -> None:
+        with self.lock:
+            self.active_requests += 1
+            self.request_count += 1
+            self.max_active_requests = max(self.max_active_requests, self.active_requests)
+            if self.active_requests == 2:
+                self.two_requests_started.set()
+        self.two_requests_started.wait(timeout=2)
+
+    def end_tracked_request(self) -> None:
+        with self.lock:
+            self.active_requests -= 1
+
+
+class _ControlledUpstreamHandler(BaseHTTPRequestHandler):
+    protocol_version = "HTTP/1.1"
+
+    def do_POST(self) -> None:
+        upstream: Final = self.server
+        assert isinstance(upstream, _ControlledUpstream)
+        length: Final = int(self.headers.get("content-length") or "0")
+        self.rfile.read(length)
+        upstream.start_request()
+        try:
+            body: Final = b"{}"
+            self.send_response(200)
+            self.send_header("content-type", "application/json")
+            self.send_header("content-length", str(len(body)))
+            self.end_headers()
+            self.wfile.write(body)
+        finally:
+            upstream.end_tracked_request()
+
+    def log_message(self, format: str, *args: object) -> None:
+        return
+
+
+@contextmanager
+def _controlled_upstream() -> Generator[_ControlledUpstream]:
+    server: Final = _ControlledUpstream()
+    thread: Final = threading.Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    try:
+        yield server
+    finally:
+        server.shutdown()
+        server.server_close()
+        thread.join(timeout=5)
+
+
+def _case(identifier: str) -> MistralOcrParityInput:
+    return MistralOcrParityInput(
+        model="mistral/mistral-ocr-latest",
+        document=ImageUrlDocument(type="image_url", image_url="https://example.com/image.png"),
+        id=identifier,
+    )
+
+
+def _sdk_call(**kwargs: object) -> object:
+    api_base: Final = cast(str, kwargs["api_base"])
+    return httpx.post(f"{api_base}/v1/ocr", content=b"{}", timeout=5)
+
+
+def test_record_cases_deduplicates_and_limits_concurrency(tmp_path: Path) -> None:
+    case_inputs: Final = (_case("one"), _case("two"), _case("one"), _case("three"))
+    with _controlled_upstream() as upstream:
+        spec: Final = ProviderSpec(model="mistral/mistral-ocr-latest", upstream_base=upstream.url, api_key="test-key")
+        results: Final = record_cases(spec, tmp_path, case_inputs, _sdk_call, max_concurrency=2)
+
+    assert len(results) == 3
+    assert upstream.request_count == 3
+    assert upstream.max_active_requests == 2