mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
ae6b762190
* refactor(tracing): normalize agent spans in Rust * refactor(tracing): generate dashboard trace types from API * test(tracing): use complete trace response fixtures * fix(ui): expose generated span error response type * fix(tracing): retain full tool call payloads * fix(tracing): preserve decoded attribute tuple shape * fix(tracing): type consumed attributes as tuple
527 lines
21 KiB
Python
527 lines
21 KiB
Python
"""
|
|
Tests for the agent tracing endpoints (litellm/proxy/tracing_endpoints.py).
|
|
"""
|
|
|
|
from collections.abc import AsyncGenerator
|
|
from contextlib import asynccontextmanager
|
|
from typing import Final
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.testclient import TestClient
|
|
|
|
from litellm.proxy import tracing_endpoints
|
|
from litellm.proxy._types import LitellmUserRoles, ProxyLifespanState, UserAPIKeyAuth
|
|
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
|
from litellm.proxy.tracing_runtime import manage_tracing, provide_storage
|
|
from litellm.rust_bridge.traces import ClickHouseStorage
|
|
from litellm.tracing import TraceReceiver, TracingPayloadTooLargeError
|
|
from litellm.tracing.store import TraceStore
|
|
from litellm.tracing.types import TraceScope
|
|
|
|
TEAM_KEY = UserAPIKeyAuth(
|
|
token="hashed-key", team_id="team-research", org_id="org-1", user_role=LitellmUserRoles.INTERNAL_USER
|
|
)
|
|
TRACE_RESPONSE: Final = {
|
|
"summary": {
|
|
"trace_id": "t1",
|
|
"name": "trace",
|
|
"service": "test",
|
|
"input_preview": "",
|
|
"start_time": "2026-01-01T00:00:00Z",
|
|
"duration_ms": 0,
|
|
"status": "ok",
|
|
"span_count": 0,
|
|
"agent_count": 0,
|
|
"agent_invocations": 0,
|
|
"llm_calls": 0,
|
|
"tool_calls": 0,
|
|
"error_count": 0,
|
|
"input_tokens": 0,
|
|
"output_tokens": 0,
|
|
"models": [],
|
|
"spend": None,
|
|
},
|
|
"agents": [],
|
|
"spans": [],
|
|
}
|
|
SPAN_DETAIL_RESPONSE: Final = {
|
|
"span_id": "s1",
|
|
"input": "",
|
|
"output": "",
|
|
"input_ui": {"kind": "text", "text": ""},
|
|
"output_ui": {"kind": "text", "text": ""},
|
|
"attributes": {},
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("auth", "scope", "can_write"),
|
|
(
|
|
pytest.param(
|
|
UserAPIKeyAuth(token="admin-key", team_id="team-a", user_role=LitellmUserRoles.PROXY_ADMIN),
|
|
TraceScope(team_ids=(), api_key_hash=""),
|
|
True,
|
|
id="admin",
|
|
),
|
|
pytest.param(
|
|
UserAPIKeyAuth(token="view-key", team_id="team-a", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY),
|
|
TraceScope(team_ids=(), api_key_hash=""),
|
|
False,
|
|
id="view-only-admin",
|
|
),
|
|
pytest.param(
|
|
TEAM_KEY,
|
|
TraceScope(team_ids=("team-research",), api_key_hash=""),
|
|
True,
|
|
id="team-key",
|
|
),
|
|
pytest.param(
|
|
UserAPIKeyAuth(token="hashed-key", user_role=LitellmUserRoles.INTERNAL_USER),
|
|
TraceScope(team_ids=("",), api_key_hash="hashed-key"),
|
|
True,
|
|
id="teamless-key",
|
|
),
|
|
),
|
|
)
|
|
def test_trace_read_and_write_permissions(
|
|
client: TestClient, receiver: MagicMock, auth: UserAPIKeyAuth, scope: TraceScope, can_write: bool
|
|
) -> None:
|
|
client.app.dependency_overrides[user_api_key_auth] = lambda: auth
|
|
|
|
read: Final = client.get("/v1/traces?start_ms=1&end_ms=2")
|
|
assert read.status_code == 200, read.text
|
|
receiver.list_traces.assert_awaited_once_with(scope=scope, start_ms=1, end_ms=2, cursor=None)
|
|
|
|
write: Final = client.post("/v1/traces", json={})
|
|
assert write.status_code == (200 if can_write else 403), write.text
|
|
if not can_write:
|
|
receiver.ingest.assert_not_awaited()
|
|
return
|
|
receiver.ingest.assert_awaited_once()
|
|
tenant: Final = receiver.ingest.await_args.kwargs["tenant"]
|
|
assert (tenant.team_id, tenant.api_key_hash, tenant.org_id) == (
|
|
auth.team_id or "",
|
|
auth.token or "",
|
|
auth.org_id or "",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def receiver(client) -> MagicMock:
|
|
fake = MagicMock()
|
|
fake.ingest = AsyncMock(return_value=1)
|
|
fake.list_traces = AsyncMock(return_value={"data": [], "next_cursor": None})
|
|
fake.get_trace = AsyncMock(return_value=None)
|
|
fake.get_span = AsyncMock(return_value=None)
|
|
client.app.dependency_overrides[tracing_endpoints.provide_receiver] = lambda: fake
|
|
return fake
|
|
|
|
|
|
@pytest.fixture
|
|
def client() -> TestClient:
|
|
app = FastAPI()
|
|
app.include_router(tracing_endpoints.router)
|
|
app.dependency_overrides[user_api_key_auth] = lambda: TEAM_KEY
|
|
return TestClient(app)
|
|
|
|
|
|
@pytest.mark.parametrize("native_available", [True, False])
|
|
def test_501_when_tracing_not_enabled(
|
|
client: TestClient, native_available: bool, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
from google.rpc.status_pb2 import Status
|
|
|
|
from litellm.rust_bridge import loader
|
|
|
|
if not native_available:
|
|
monkeypatch.setattr(loader, "_cached_bridge", None)
|
|
response: Final = client.post("/v1/traces", content=b"")
|
|
assert response.status_code == 501
|
|
assert response.headers["content-type"] == "application/x-protobuf"
|
|
assert Status.FromString(response.content).message == (
|
|
"Agent tracing is not enabled. Set `tracing:` in general_settings and CLICKHOUSE_URL." if native_available else ""
|
|
)
|
|
assert client.get("/v1/traces").status_code == 501
|
|
|
|
|
|
def test_post_protobuf_returns_empty_protobuf(client, receiver):
|
|
response = client.post(
|
|
"/v1/traces",
|
|
content=b"\x0a\x00",
|
|
headers={"content-type": "application/x-protobuf", "content-encoding": "gzip"},
|
|
)
|
|
assert response.status_code == 200
|
|
assert response.content == b""
|
|
assert response.headers["content-type"] == "application/x-protobuf"
|
|
kwargs = receiver.ingest.call_args.kwargs
|
|
assert kwargs["body"] is not None
|
|
assert kwargs["content_type"] == "application/x-protobuf"
|
|
assert kwargs["content_encoding"] == "gzip"
|
|
assert kwargs["tenant"].team_id == "team-research"
|
|
|
|
|
|
def test_post_json_returns_empty_json(client, receiver):
|
|
response = client.post("/v1/traces", content=b"{}", headers={"content-type": "application/json"})
|
|
assert response.status_code == 200
|
|
assert response.json() == {}
|
|
|
|
|
|
def test_post_clickhouse_failure_is_503_with_retry_after(client, receiver):
|
|
receiver.ingest.side_effect = RuntimeError("ClickHouse unavailable")
|
|
response = client.post("/v1/traces", content=b"", headers={"content-type": "application/x-protobuf"})
|
|
assert response.status_code == 503
|
|
assert response.headers["retry-after"] == str(tracing_endpoints.OTLP_RETRY_AFTER_SECONDS)
|
|
|
|
|
|
def test_post_too_large_is_413(client, receiver):
|
|
receiver.ingest.side_effect = TracingPayloadTooLargeError("OTLP body exceeds 10 bytes")
|
|
response = client.post("/v1/traces", content=b"x" * 20)
|
|
assert response.status_code == 413
|
|
from google.rpc.status_pb2 import Status
|
|
|
|
assert "exceeds" in Status.FromString(response.content).message
|
|
|
|
|
|
def test_list_traces_passes_scope_window_and_cursor(client, receiver):
|
|
response = client.get("/v1/traces", params={"start_ms": 1, "end_ms": 2, "cursor": "abc"})
|
|
assert response.status_code == 200
|
|
assert response.json() == {"data": [], "next_cursor": None}
|
|
receiver.list_traces.assert_awaited_once_with(
|
|
scope={"team_ids": ("team-research",), "api_key_hash": ""}, start_ms=1, end_ms=2, cursor="abc"
|
|
)
|
|
|
|
|
|
def test_list_traces_defaults_to_last_24h(client, receiver):
|
|
client.get("/v1/traces")
|
|
kwargs = receiver.list_traces.call_args.kwargs
|
|
assert kwargs["end_ms"] - kwargs["start_ms"] == tracing_endpoints.MS_PER_DAY
|
|
assert kwargs["cursor"] is None
|
|
|
|
|
|
def test_get_trace_404_and_200(client, receiver):
|
|
assert client.get("/v1/traces/missing").status_code == 404
|
|
receiver.get_trace.return_value = TRACE_RESPONSE
|
|
response = client.get("/v1/traces/t1")
|
|
assert response.status_code == 200
|
|
assert response.json() == TRACE_RESPONSE
|
|
receiver.get_trace.assert_awaited_with("t1", {"team_ids": ("team-research",), "api_key_hash": ""}, "")
|
|
|
|
|
|
def test_get_span_404_and_200(client, receiver):
|
|
assert client.get("/v1/traces/t1/spans/s1").status_code == 404
|
|
receiver.get_span.return_value = SPAN_DETAIL_RESPONSE
|
|
response = client.get("/v1/traces/t1/spans/s1")
|
|
assert response.status_code == 200
|
|
assert response.json()["span_id"] == "s1"
|
|
receiver.get_span.assert_awaited_with("t1", "s1", {"team_ids": ("team-research",), "api_key_hash": ""}, "")
|
|
|
|
|
|
def test_get_span_serves_ui_content_from_stored_payloads(client):
|
|
storage = MagicMock()
|
|
stored_output = '{"role": "ai", "content": "", "tool_calls": [{"name": "lookup", "args": {"id": 7}}]}'
|
|
storage.query = AsyncMock(
|
|
return_value=[{"span_id": "s1", "input": '{"city": "Paris"}', "output": stored_output, "attributes": {}}]
|
|
)
|
|
client.app.dependency_overrides[tracing_endpoints.provide_receiver] = lambda: TraceReceiver(TraceStore(storage))
|
|
body = client.get("/v1/traces/t1/spans/s1").json()
|
|
assert body["output"] == stored_output
|
|
assert body["input_ui"] == {"kind": "fields", "fields": [{"key": "city", "value": "Paris"}]}
|
|
assert body["output_ui"] == {
|
|
"kind": "messages",
|
|
"messages": [
|
|
{"role": "assistant", "content": "", "tool_calls": [{"name": "lookup", "arguments": '{"id": 7}'}]}
|
|
],
|
|
}
|
|
|
|
|
|
def test_trace_detail_passes_scoped_reference(client, receiver):
|
|
receiver.get_trace.return_value = TRACE_RESPONSE
|
|
assert client.get("/v1/traces/t1?trace_ref=run-one").status_code == 200
|
|
receiver.get_trace.assert_awaited_with("t1", {"team_ids": ("team-research",), "api_key_hash": ""}, "run-one")
|
|
|
|
|
|
def test_invalid_export_and_cursor_are_client_errors(client, receiver):
|
|
from litellm.tracing.decode import InvalidOTLPPayloadError
|
|
|
|
receiver.ingest.side_effect = InvalidOTLPPayloadError("invalid OTLP trace payload")
|
|
assert client.post("/v1/traces", content=b"broken").status_code == 400
|
|
receiver.list_traces.side_effect = ValueError("Invalid trace cursor")
|
|
assert client.get("/v1/traces?cursor=broken").status_code == 400
|
|
|
|
|
|
def test_teamless_key_without_token_gets_403_on_reads(client, receiver):
|
|
client.app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
|
user_role=LitellmUserRoles.INTERNAL_USER
|
|
)
|
|
assert client.get("/v1/traces").status_code == 403
|
|
receiver.list_traces.assert_not_called()
|
|
|
|
|
|
def test_view_only_admin_cannot_ingest_traces(client, receiver):
|
|
client.app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
|
token="admin-key", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
|
|
)
|
|
response = client.post("/v1/traces", content=b"{}")
|
|
assert response.status_code == 403
|
|
receiver.ingest.assert_not_called()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"status_code, field, message",
|
|
[(401, "detail", "Invalid API key"), (403, "message", "Not allowed to ingest agent traces")],
|
|
)
|
|
def test_auth_failure_precedes_disabled_receiver(
|
|
client: TestClient, status_code: int, field: str, message: str
|
|
) -> None:
|
|
def unavailable() -> None:
|
|
return None
|
|
|
|
def authenticate() -> UserAPIKeyAuth:
|
|
if status_code == 401:
|
|
raise HTTPException(status_code=401, detail="Invalid API key")
|
|
return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)
|
|
|
|
client.app.dependency_overrides[user_api_key_auth] = authenticate
|
|
client.app.dependency_overrides[tracing_endpoints.provide_receiver] = unavailable
|
|
response: Final = client.post("/v1/traces", content=b"{}", headers={"content-type": "application/json"})
|
|
assert response.status_code == status_code
|
|
assert response.json() == {field: message}
|
|
|
|
|
|
def test_disabled_receiver_precedes_read_scope_rejection(client: TestClient) -> None:
|
|
client.app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
|
user_role=LitellmUserRoles.INTERNAL_USER
|
|
)
|
|
response: Final = client.get("/v1/traces")
|
|
assert response.status_code == 501
|
|
assert response.json() == {
|
|
"detail": "Agent tracing is not enabled. Set `tracing:` in general_settings and CLICKHOUSE_URL."
|
|
}
|
|
|
|
|
|
@pytest.mark.requires_rust_extension
|
|
def test_injected_receiver_persists_authenticated_tenant(client: TestClient) -> None:
|
|
storage: Final = MagicMock(spec=ClickHouseStorage)
|
|
storage.insert_rows = AsyncMock()
|
|
tracing: Final = TraceReceiver(TraceStore(storage))
|
|
client.app.dependency_overrides[tracing_endpoints.provide_receiver] = lambda: tracing
|
|
response: Final = client.post(
|
|
"/v1/traces",
|
|
json={
|
|
"resourceSpans": [
|
|
{
|
|
"resource": {
|
|
"attributes": [
|
|
{"key": "litellm.team_id", "value": {"stringValue": "spoofed-team"}},
|
|
{"key": "litellm.api_key_hash", "value": {"stringValue": "spoofed-key"}},
|
|
{"key": "litellm.org_id", "value": {"stringValue": "spoofed-org"}},
|
|
]
|
|
},
|
|
"scopeSpans": [
|
|
{
|
|
"spans": [
|
|
{
|
|
"traceId": "01" * 16,
|
|
"spanId": "02" * 8,
|
|
"name": "dependency-injection",
|
|
"startTimeUnixNano": "1000000000",
|
|
"endTimeUnixNano": "1000000001",
|
|
}
|
|
]
|
|
}
|
|
],
|
|
}
|
|
],
|
|
},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
assert response.json() == {}
|
|
storage.insert_rows.assert_awaited_once()
|
|
table, rows = storage.insert_rows.await_args.args
|
|
assert table == "otel_traces"
|
|
assert len(rows) == 1
|
|
assert rows[0]["TeamId"] == TEAM_KEY.team_id
|
|
assert rows[0]["ApiKeyHash"] == TEAM_KEY.token
|
|
assert rows[0]["ResourceAttributes"] == {
|
|
"litellm.team_id": TEAM_KEY.team_id,
|
|
"litellm.api_key_hash": TEAM_KEY.token,
|
|
"litellm.org_id": TEAM_KEY.org_id,
|
|
}
|
|
|
|
|
|
def test_lifespan_receivers_are_app_local() -> None:
|
|
first_storage: Final = MagicMock(spec=ClickHouseStorage)
|
|
first_storage.query = AsyncMock(
|
|
return_value=[
|
|
{
|
|
"span_id": "first-span",
|
|
"input": "first-input",
|
|
"output": "",
|
|
"attributes": {},
|
|
}
|
|
]
|
|
)
|
|
second_storage: Final = MagicMock(spec=ClickHouseStorage)
|
|
second_storage.query = AsyncMock(
|
|
return_value=[
|
|
{
|
|
"span_id": "second-span",
|
|
"input": "second-input",
|
|
"output": "",
|
|
"attributes": {},
|
|
}
|
|
]
|
|
)
|
|
first_receiver: Final = TraceReceiver(TraceStore(first_storage))
|
|
second_receiver: Final = TraceReceiver(TraceStore(second_storage))
|
|
first_storage.ensure_schema = AsyncMock()
|
|
second_storage.ensure_schema = AsyncMock()
|
|
|
|
@asynccontextmanager
|
|
async def first_lifespan(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]:
|
|
async with manage_tracing(True, lambda: first_receiver) as receiver:
|
|
state: Final[ProxyLifespanState] = {"tracing_receiver": receiver}
|
|
yield state
|
|
|
|
@asynccontextmanager
|
|
async def second_lifespan(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]:
|
|
async with manage_tracing(True, lambda: second_receiver) as receiver:
|
|
state: Final[ProxyLifespanState] = {"tracing_receiver": receiver}
|
|
yield state
|
|
|
|
first_app: Final = FastAPI(lifespan=first_lifespan)
|
|
second_app: Final = FastAPI(lifespan=second_lifespan)
|
|
first_app.include_router(tracing_endpoints.router)
|
|
second_app.include_router(tracing_endpoints.router)
|
|
first_app.dependency_overrides[user_api_key_auth] = lambda: TEAM_KEY
|
|
second_app.dependency_overrides[user_api_key_auth] = lambda: TEAM_KEY
|
|
|
|
with TestClient(first_app) as first_client:
|
|
with TestClient(second_app) as second_client:
|
|
second_response: Final = second_client.get("/v1/traces/t1/spans/second-span?trace_ref=second-run")
|
|
simultaneous: Final = first_client.get("/v1/traces/t1/spans/first-span?trace_ref=first-run")
|
|
first_response: Final = first_client.get("/v1/traces/t1/spans/first-span?trace_ref=first-run")
|
|
assert simultaneous.json() == first_response.json()
|
|
first_storage.ensure_schema.assert_awaited_once()
|
|
second_storage.ensure_schema.assert_awaited_once()
|
|
|
|
assert first_response.status_code == second_response.status_code == 200
|
|
assert first_response.json() == {
|
|
"span_id": "first-span",
|
|
"input": "first-input",
|
|
"output": "",
|
|
"attributes": {},
|
|
"input_ui": {"kind": "text", "text": "first-input"},
|
|
"output_ui": {"kind": "text", "text": ""},
|
|
}
|
|
assert second_response.json() == {
|
|
"span_id": "second-span",
|
|
"input": "second-input",
|
|
"output": "",
|
|
"attributes": {},
|
|
"input_ui": {"kind": "text", "text": "second-input"},
|
|
"output_ui": {"kind": "text", "text": ""},
|
|
}
|
|
assert first_storage.query.await_count == 2
|
|
first_storage.query.assert_awaited_with(
|
|
"span_detail",
|
|
{
|
|
"team_ids": (TEAM_KEY.team_id,),
|
|
"api_key_hash": "",
|
|
"trace_id": "t1",
|
|
"span_id": "first-span",
|
|
"trace_ref": "first-run",
|
|
},
|
|
)
|
|
second_storage.query.assert_awaited_once_with(
|
|
"span_detail",
|
|
{
|
|
"team_ids": (TEAM_KEY.team_id,),
|
|
"api_key_hash": "",
|
|
"trace_id": "t1",
|
|
"span_id": "second-span",
|
|
"trace_ref": "second-run",
|
|
},
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("auth", [TEAM_KEY, UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER)])
|
|
def test_query_validation_precedes_trace_access_checks(client: TestClient, auth: UserAPIKeyAuth) -> None:
|
|
client.app.dependency_overrides[user_api_key_auth] = lambda: auth
|
|
response: Final = client.get("/v1/traces", params={"start_ms": "invalid"})
|
|
assert response.status_code == 422
|
|
assert response.json()["detail"][0]["loc"] == ["query", "start_ms"]
|
|
|
|
|
|
@pytest.mark.parametrize("enabled", [True, False])
|
|
def test_unavailable_lifespan_receiver_returns_501(enabled: bool) -> None:
|
|
storage: Final = MagicMock(spec=ClickHouseStorage)
|
|
storage.ensure_schema = AsyncMock(side_effect=RuntimeError("storage unavailable"))
|
|
tracing: Final = TraceReceiver(TraceStore(storage))
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]:
|
|
async with manage_tracing(enabled, lambda: tracing) as receiver:
|
|
state: Final[ProxyLifespanState] = {"tracing_receiver": receiver}
|
|
yield state
|
|
|
|
app: Final = FastAPI(lifespan=lifespan)
|
|
app.include_router(tracing_endpoints.router)
|
|
app.dependency_overrides[user_api_key_auth] = lambda: TEAM_KEY
|
|
with TestClient(app) as client:
|
|
response: Final = client.get("/v1/traces")
|
|
assert response.status_code == 501
|
|
assert storage.ensure_schema.await_count == int(enabled)
|
|
storage.query.assert_not_called()
|
|
|
|
|
|
def test_lens_reads_from_the_lifespan_storage() -> None:
|
|
from litellm.proxy.lens.endpoints import router as lens_router
|
|
|
|
storage: Final = MagicMock(spec=ClickHouseStorage)
|
|
storage.ensure_schema = AsyncMock()
|
|
storage.lens_sample = AsyncMock(return_value=[])
|
|
tracing: Final = TraceReceiver(TraceStore(storage))
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]:
|
|
async with manage_tracing(True, lambda: tracing) as receiver:
|
|
state: Final[ProxyLifespanState] = {"tracing_receiver": receiver}
|
|
yield state
|
|
|
|
app: Final = FastAPI(lifespan=lifespan)
|
|
app.include_router(lens_router)
|
|
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
|
|
with TestClient(app) as client:
|
|
response: Final = client.post(
|
|
"/lens/preview/sample",
|
|
json={"settings": {"name": "Review", "model": "analysis", "context": "Find failed executions"}},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
assert response.json()["executions"] == []
|
|
storage.lens_sample.assert_awaited_once()
|
|
assert storage.lens_sample.await_args.args[0]["all_teams"] == 1
|
|
|
|
|
|
def test_lens_reads_from_injected_storage_without_receiver() -> None:
|
|
from litellm.proxy.lens.endpoints import router as lens_router
|
|
from litellm.proxy.lens.sources import Storage
|
|
|
|
storage: Final = MagicMock(spec=Storage)
|
|
storage.lens_sample = AsyncMock(return_value=[])
|
|
app: Final = FastAPI()
|
|
app.include_router(lens_router)
|
|
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
|
|
app.dependency_overrides[provide_storage] = lambda: storage
|
|
|
|
with TestClient(app) as client:
|
|
response: Final = client.post(
|
|
"/lens/preview/sample",
|
|
json={"settings": {"name": "Review", "model": "analysis", "context": "Find failed executions"}},
|
|
)
|
|
|
|
assert response.status_code == 200, response.text
|
|
assert response.json()["executions"] == []
|
|
storage.lens_sample.assert_awaited_once()
|