litellm/tests/rust-python-harness/strategies/trace_parity/test_runner.py
Yujong Lee 96baeb8b04 refactor(rust): remove gateway, config, router, realtime, and trace-parity infrastructure
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-09-16 16:00:07 +00:00

242 lines
10 KiB
Python

from __future__ import annotations
import importlib
import os
from pathlib import Path
from types import SimpleNamespace
from typing import Final, cast
import pytest
import litellm
from ...shared.reporting.models import Coverage, HarnessCase, HarnessRun, RunStatus, SdkFunction, Surface
from ...shared.reporting.strategy import ModuleCaseSpec
from ...shared.tracing.profiler import FunctionTraceEvent
from ...shared.tracing.steps import PipelineStep
from .models import RouteFixture, RouteSpec, TraceScenario, TraceSuite
from .reporting import TraceArtifact
from .runner import run_trace_cases, run_trace_scenario, scenario_nodeids, validate_trace_suite
from .sdk.execution import SdkCall, collect_trace, execute_trace
def _fixture(_base_url: str) -> RouteFixture:
return RouteFixture(kwargs={}, provider_responses=())
def _case(*, surface: Surface = "sdk", function: SdkFunction = "ocr") -> HarnessCase:
return HarnessCase(
strategy_id="trace_parity",
strategy_label="Trace parity",
sdk_function=function,
spec=ModuleCaseSpec(coverage=Coverage.PARTIAL, module="example"),
surface=surface,
)
def test_scenario_filtering_and_occurrence_node_ids() -> None:
suite: Final = TraceSuite(
route=RouteSpec("ocr", ("ocr", "aocr"), _fixture),
scenarios=(
TraceScenario("sync-one", _fixture, asynchronous=False),
TraceScenario("async-one", _fixture, asynchronous=True),
TraceScenario("async-two", _fixture, asynchronous=True),
),
)
case: Final = _case()
nodes: Final = scenario_nodeids(suite, case, frozenset({"async-two"}))
assert tuple(nodeid for _, nodeid in nodes) == ("trace:sdk:ocr:async-two",)
def test_runner_arguments_select_scenarios(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
runner: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.runner")
case: Final = _case()
selected: list[frozenset[str]] = []
def capture_case(
_run: HarnessRun,
_case: HarnessCase,
scenarios: frozenset[str],
_on_update: object,
) -> None:
selected.append(scenarios)
monkeypatch.setattr(runner, "_run_case", capture_case)
exit_code, _ = run_trace_cases((case,), tmp_path, lambda _: None, ("mistral",))
assert exit_code == 0
assert selected == [frozenset({"mistral"})]
def test_python_trace_preserves_native_ocr_dispatch_setting(monkeypatch: pytest.MonkeyPatch) -> None:
execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.execution")
route: Final = RouteSpec("ocr", ("ocr", "aocr"), _fixture)
observed: list[str | None] = []
def collect(
_function: SdkCall,
_fixture: RouteFixture,
*,
asynchronous: bool,
) -> SimpleNamespace:
observed.append(os.environ.get("LITELLM_RUST"))
return SimpleNamespace(
events=(FunctionTraceEvent(0, None, "aocr" if asynchronous else "ocr"),),
error=None,
)
monkeypatch.setattr(execution, "_collect", collect)
monkeypatch.setenv("LITELLM_RUST", "0")
collect_trace(route, asynchronous=False)
monkeypatch.setenv("LITELLM_RUST", "1")
collect_trace(route, asynchronous=True)
assert observed == ["0", "1"]
assert os.environ["LITELLM_RUST"] == "1"
def test_expected_provider_failure_omits_feedback_banner(
capsys: pytest.CaptureFixture[str], monkeypatch: pytest.MonkeyPatch
) -> None:
loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.responses.case")
suite: Final = cast(TraceSuite, loaded.TRACE_SUITE)
scenario: Final = next(item for item in suite.scenarios if item.name == "async-openai-provider-error")
monkeypatch.setattr(litellm, "suppress_debug_info", False)
result: Final = execute_trace(suite.route, scenario, "sdk")
assert result.python_error is None
assert "Give Feedback / Get Help" not in capsys.readouterr().out
assert litellm.suppress_debug_info is False
@pytest.mark.parametrize("asynchronous", (False, True))
def test_vertex_trace_keeps_unmapped_helpers_and_parents(asynchronous: bool) -> None:
loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case")
suite: Final = cast(TraceSuite, loaded.TRACE_SUITE)
name: Final = f"{'async' if asynchronous else 'sync'}-vertex-deepseek"
scenario: Final = next(item for item in suite.scenarios if item.name == name)
trace: Final = execute_trace(suite.route, scenario, "sdk")
assert trace.python_error is None
url: Final = next(
event for event in trace.python if event.raw.endswith(" VertexAIDeepSeekOCRConfig.get_complete_url")
)
project: Final = next(
event for event in trace.python if event.raw.endswith(" VertexBase.safe_get_vertex_ai_project")
)
location: Final = next(
event for event in trace.python if event.raw.endswith(" VertexBase.safe_get_vertex_ai_location")
)
assert project.parent_id == location.parent_id == url.id
assert not any(event.raw.endswith(" VertexBase.get_access_token") for event in trace.python)
@pytest.mark.parametrize("asynchronous", (False, True))
def test_vertex_credentials_trace_runs_real_auth_helpers(asynchronous: bool, monkeypatch: pytest.MonkeyPatch) -> None:
loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case")
suite: Final = cast(TraceSuite, loaded.TRACE_SUITE)
name: Final = f"{'async' if asynchronous else 'sync'}-vertex-deepseek-credentials"
scenario: Final = next(item for item in suite.scenarios if item.name == name)
monkeypatch.setenv("VERTEXAI_CREDENTIALS", "original-credentials")
monkeypatch.setenv("VERTEX_AI_API_KEY", "original-api-key")
trace: Final = execute_trace(suite.route, scenario, "sdk")
assert trace.python_error is None
validate: Final = next(
event for event in trace.python if event.raw.endswith(" VertexAIDeepSeekOCRConfig.validate_environment")
)
helpers: Final = (
"VertexBase.safe_get_vertex_ai_project",
"VertexBase.safe_get_vertex_ai_credentials",
"VertexBase.get_access_token",
)
assert tuple(event.raw.split(" ", 1)[1] for event in trace.python if event.parent_id == validate.id) == helpers
token: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.get_access_token"))
load: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.load_auth"))
refresh: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.refresh_auth"))
assert load.parent_id == token.id
assert refresh.parent_id == load.id
assert os.environ["VERTEXAI_CREDENTIALS"] == "original-credentials"
assert os.environ["VERTEX_AI_API_KEY"] == "original-api-key"
def test_scenario_validation_rejects_duplicate_and_unsafe_names() -> None:
route: Final = RouteSpec("ocr", ("ocr", "aocr"), _fixture)
duplicate: Final = TraceSuite(
route=route,
scenarios=(
TraceScenario("sync-same", _fixture, asynchronous=False),
TraceScenario("sync-same", _fixture, asynchronous=False),
),
)
unsafe: Final = TraceSuite(route=route, scenarios=(TraceScenario("sync-bad:name", _fixture, asynchronous=False),))
case: Final = _case()
assert validate_trace_suite(duplicate, case) is not None
assert validate_trace_suite(unsafe, case) is not None
def test_scenario_validation_rejects_invalid_names_and_route_registration() -> None:
invalid_name: Final = TraceSuite(
route=RouteSpec("ocr", ("ocr", "aocr"), _fixture),
scenarios=(TraceScenario("bedrock", _fixture, asynchronous=True),),
)
wrong_function: Final = TraceSuite(
route=RouteSpec("messages", ("create", "acreate"), _fixture),
scenarios=(TraceScenario("sync-one", _fixture, asynchronous=False),),
)
wrong_surface: Final = TraceSuite(
route=RouteSpec("ocr", ("ocr", "aocr"), _fixture),
scenarios=(TraceScenario("sync-one", _fixture, asynchronous=False),),
)
case: Final = _case()
assert "start with sync- or async-" in (validate_trace_suite(invalid_name, case) or "")
assert "does not match case function" in (validate_trace_suite(wrong_function, case) or "")
assert "requires the sdk surface" in (validate_trace_suite(wrong_surface, _case(surface="gateway")) or "")
def test_invalid_route_dispatch_records_harness_error() -> None:
case: Final = _case()
run: Final = HarnessRun.from_cases((case,))
result: Final = run.results[case.key]
suite: Final = TraceSuite(
route=RouteSpec("ocr", ("ocr", "aocr"), _fixture),
scenarios=(TraceScenario("sync-one", _fixture, asynchronous=False),),
)
nodeid: Final = "trace:gateway:ocr:sync-one"
run_trace_scenario(run, result, suite, suite.scenarios[0], "gateway", nodeid, lambda _: None)
assert result.outcomes[nodeid] is RunStatus.ERROR
assert run.failures == [(nodeid, "trace scenarios only run on the sdk surface")]
def test_python_trace_without_errors_passes(monkeypatch: pytest.MonkeyPatch) -> None:
runner: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.runner")
case: Final = _case()
run: Final = HarnessRun.from_cases((case,))
result: Final = run.results[case.key]
suite: Final = TraceSuite(
route=RouteSpec("ocr", ("ocr", "aocr"), _fixture),
scenarios=(TraceScenario("sync-one", _fixture, asynchronous=False),),
)
trace: Final = TraceArtifact.from_traces(
surface="sdk",
sdk_function="ocr",
scenario="sync-one",
python=(PipelineStep(0, None, "python_step", "python.py:1 python_step"),),
)
monkeypatch.setattr(runner, "_execute_scenario", lambda *_args: trace)
run_trace_scenario(run, result, suite, suite.scenarios[0], "sdk", "trace:sdk:ocr:sync-one", lambda _: None)
assert result.outcomes["trace:sdk:ocr:sync-one"] is RunStatus.PASSED
assert run.failures == []