mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
242 lines
10 KiB
Python
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 == []
|