litellm/tests/rust-python-harness/shared/parity/fixtures/test_pipeline.py

from __future__ import annotations

import logging
import threading
from contextlib import AbstractContextManager
from dataclasses import dataclass
from pathlib import Path
from typing import Final, Literal

import httpx
import pytest
from hypothesis import strategies as st
from pydantic import BaseModel, ConfigDict

from ..local_server import LocalHttpHandler, LocalHttpServer, serve_in_thread
from ..recorded_http import RecordedResponse
from .pipeline import (
    RecordingInvocation,
    RecordingTarget,
    build_recording_jobs,
    record_fixtures,
)
from .recording import UpstreamEndpoint
from .store import fixture_path


class _FixtureInput(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    identifier: str

    def canonical_input(self) -> dict[str, object]:
        return {"identifier": self.identifier}


class _ParityCase(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    litellm_input: _FixtureInput
    provider_responses: tuple[RecordedResponse, ...]


class _Upstream(LocalHttpServer):
    def __init__(self, status: int = 200) -> None:
        super().__init__(("127.0.0.1", 0), _UpstreamHandler)
        self.response_status: Final = status

class _UpstreamHandler(LocalHttpHandler):

    def do_POST(self) -> None:
        length: Final = int(self.headers.get("content-length") or "0")
        self.rfile.read(length)
        body: Final = b"{}"
        server: Final = self.server
        assert isinstance(server, _Upstream)
        self.send_response(server.response_status)
        self.send_header("content-type", "application/json")
        self.send_header("content-length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

def _upstream(status: int = 200) -> AbstractContextManager[_Upstream]:
    return serve_in_thread(_Upstream(status))


@dataclass(frozen=True, slots=True)
class _OrderedInvocation:
    order: Literal["slow", "fast"]
    slow_started: threading.Event
    fast_finished: threading.Event

    def execute(self, provider_url: str, case_input: _FixtureInput) -> None:
        if self.order == "slow":
            self.slow_started.set()
            if not self.fast_finished.wait(timeout=2):
                raise TimeoutError("fast recording did not finish")
        else:
            if not self.slow_started.wait(timeout=2):
                raise TimeoutError("slow recording did not start")
        response: Final = httpx.post(f"{provider_url}/record", json={"id": case_input.identifier}, timeout=5)
        response.raise_for_status()
        if self.order == "fast":
            self.fast_finished.set()


@dataclass(frozen=True, slots=True)
class _Invocation:
    def execute(self, provider_url: str, case_input: _FixtureInput) -> None:
        response: Final = httpx.post(f"{provider_url}/record", json={"id": case_input.identifier}, timeout=5)
        response.raise_for_status()


def _target(
    name: str,
    upstream_url: str,
    case_input: _FixtureInput,
    invocation: RecordingInvocation[_FixtureInput],
) -> RecordingTarget[_FixtureInput]:
    return RecordingTarget(
        name=name,
        upstream=UpstreamEndpoint(base_url=upstream_url),
        strategy=st.just(case_input),
        invocation=invocation,
        required_inputs=(case_input,),
    )


def test_build_jobs_keeps_required_inputs_before_generated_inputs_and_deduplicates(tmp_path: Path) -> None:
    required: Final = _FixtureInput(identifier="required")
    generated: Final = _FixtureInput(identifier="generated")
    target: Final = RecordingTarget(
        name="ordered",
        upstream=UpstreamEndpoint(base_url="https://provider.invalid"),
        strategy=st.just(generated),
        invocation=_Invocation(),
        required_inputs=(required, required),
    )

    jobs: Final = build_recording_jobs((target,), tmp_path, examples=1)

    assert tuple(job.case_input.identifier for job in jobs) == ("required", "generated")


def test_progress_follows_completion_order(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
    slow_started: Final = threading.Event()
    fast_finished: Final = threading.Event()
    with _upstream() as upstream:
        targets: Final = (
            _target(
                "slow",
                upstream.url,
                _FixtureInput(identifier="slow"),
                _OrderedInvocation("slow", slow_started, fast_finished),
            ),
            _target(
                "fast",
                upstream.url,
                _FixtureInput(identifier="fast"),
                _OrderedInvocation("fast", slow_started, fast_finished),
            ),
        )
        with caplog.at_level(logging.INFO, logger="tests.rust-python-harness.shared.parity.fixtures.pipeline"):
            summary: Final = record_fixtures(targets, tmp_path, 1, 2, _ParityCase)

    progress: Final = tuple(record.message for record in caplog.records if record.message.startswith("["))
    assert len(summary.recorded) == 2
    assert summary.exit_code == 0
    assert "recorded fast" in progress[0]
    assert "recorded slow" in progress[1]
    assert caplog.records[0].message == "Recording 2 fixtures across 2 targets with concurrency 2"
    assert caplog.records[-1].message == "Finished 2 fixtures: 2 recorded, 0 cached, 0 failed"


def test_failure_does_not_stop_independent_recordings(tmp_path: Path) -> None:
    stale_input: Final = _FixtureInput(identifier="stale")
    stale_directory: Final = tmp_path / "stale"
    stale_directory.mkdir()
    fixture_path(stale_directory, stale_input).with_suffix(".json").write_text(
        '{"schema_version": 0}\n', encoding="utf-8"
    )
    with _upstream() as upstream:
        targets: Final = (
            _target("stale", upstream.url, stale_input, _Invocation()),
            _target("valid", upstream.url, _FixtureInput(identifier="valid"), _Invocation()),
        )
        summary: Final = record_fixtures(targets, tmp_path, 1, 2, _ParityCase)

    assert len(summary.recorded) == 1
    assert summary.recorded[0].target_name == "valid"
    assert len(summary.failed) == 1
    assert summary.failed[0].target_name == "stale"
    assert summary.exit_code == 1


@pytest.mark.parametrize("status", (408, 429, 500, 503))
def test_transient_response_is_not_cached_and_can_be_retried(tmp_path: Path, status: int) -> None:
    case_input: Final = _FixtureInput(identifier="retry")
    with _upstream(status) as upstream:
        target: Final = _target("retry", upstream.url, case_input, _Invocation())
        failed: Final = record_fixtures((target,), tmp_path, 1, 1, _ParityCase)
    assert failed.exit_code == 1
    assert not fixture_path(tmp_path / "retry", case_input).exists()
    with _upstream() as healthy_upstream:
        healthy_target: Final = _target("retry", healthy_upstream.url, case_input, _Invocation())
        retried: Final = record_fixtures((healthy_target,), tmp_path, 1, 1, _ParityCase)
    assert retried.exit_code == 0
    assert len(retried.recorded) == 1


def test_provider_rejected_response_can_be_recorded(tmp_path: Path) -> None:
    with _upstream(400) as upstream:
        target: Final = _target("rejected", upstream.url, _FixtureInput(identifier="invalid"), _Invocation())
        summary: Final = record_fixtures((target,), tmp_path, 1, 1, _ParityCase)
    assert summary.exit_code == 0
    assert len(summary.recorded) == 1