litellm/tests/rust-python-harness/shared/parity/inprocess.py
yujonglee 2c30fe16b0
Merge pull request #38765 from BerriAI/litellm_ocr_sdk_parity_tests
test(harness): add OCR parity with migration strategy runners
2026-09-03 10:16:35 -07:00

47 lines
1.5 KiB
Python

from __future__ import annotations
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import Final, Generic, TypeVar
from .models import CapturedRequest
from .recorded_http import RecordedResponse
from .replay import ReplayServer
ResponseT = TypeVar("ResponseT")
@dataclass(frozen=True, slots=True)
class InProcessExecution(Generic[ResponseT]):
requests: tuple[CapturedRequest, ...]
response: ResponseT
def run_in_process(
provider: ReplayServer,
recorded_responses: tuple[RecordedResponse, ...],
call: Callable[[str], ResponseT],
) -> InProcessExecution[ResponseT]:
for recorded_response in recorded_responses:
provider.enqueue_response(recorded_response)
try:
response: Final = call(provider.url)
return InProcessExecution(requests=provider.take_requests(len(recorded_responses)), response=response)
except Exception:
provider.reset()
raise
async def run_in_process_async(
provider: ReplayServer,
recorded_responses: tuple[RecordedResponse, ...],
call: Callable[[str], Awaitable[ResponseT]],
) -> InProcessExecution[ResponseT]:
for recorded_response in recorded_responses:
provider.enqueue_response(recorded_response)
try:
response: Final = await call(provider.url)
return InProcessExecution(requests=provider.take_requests(len(recorded_responses)), response=response)
except Exception:
provider.reset()
raise