test: drop the deleted otel base class import and restore mcp redaction and exact langfuse usage checks

This commit is contained in:
Yuneng Jiang 2026-10-10 02:37:52 -07:00
parent f6801b8ed9
commit 36a7b11aca
No known key found for this signature in database
7 changed files with 83 additions and 50 deletions

View file

@ -251,17 +251,11 @@ class Case:
async def run(self, credentials: Mapping[str, object]) -> OCRResponse:
kwargs: Final = {**self.provider.params, **credentials}
document: Final = self.document.build()
try:
response: Final = (
await _OCR_CLIENT.aocr(model=self.provider.model, document=document, **kwargs)
if self.call == "async"
else _OCR_CLIENT.ocr(model=self.provider.model, document=document, **kwargs)
)
except Exception as exc:
reason: Final = str(exc).casefold()
if any(value in reason for value in ("quota", "rate limit", "region")):
pytest.skip(f"{self.provider.id} rejected the account due to quota or region")
raise
response: Final = (
await _OCR_CLIENT.aocr(model=self.provider.model, document=document, **kwargs)
if self.call == "async"
else _OCR_CLIENT.ocr(model=self.provider.model, document=document, **kwargs)
)
assert isinstance(response, OCRResponse)
return response

View file

@ -21,17 +21,10 @@
"extra_body": "{}"
},
"langfuse.observation.type": "generation",
"langfuse.observation.output": {
"content": "Hello! How can I assist you today?",
"role": "assistant",
"tool_calls": null,
"function_call": null,
"provider_specific_fields": null
},
"langfuse.observation.usage_details": {
"input": 10,
"output": 20,
"total": 30,
"output": 10,
"total": 20,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0
},

View file

@ -3,7 +3,7 @@
"parent_span_id": null,
"attributes": {
"langfuse.observation.cost_details": {
"total": 3.5e-05
"total": 1.8e-05
},
"langfuse.observation.input": {
"messages": [
@ -28,9 +28,9 @@
},
"langfuse.observation.type": "generation",
"langfuse.observation.usage_details": {
"input": 10,
"output": 20,
"total": 30,
"input": 9,
"output": 9,
"total": 18,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0
},

View file

@ -32,7 +32,6 @@ _EXPORT_URL: Final = "https://langfuse.unit/api/public/otel/v1/traces"
_VOLATILE_ATTRIBUTES: Final = frozenset(
{
"langfuse.observation.completion_start_time",
"langfuse.observation.cost_details",
"langfuse.observation.metadata.applied_guardrails",
"langfuse.observation.metadata.cache_hit",
"langfuse.observation.metadata.hidden_params",
@ -41,7 +40,6 @@ _VOLATILE_ATTRIBUTES: Final = frozenset(
"langfuse.observation.metadata.requester_metadata",
"langfuse.observation.metadata.response_id",
"langfuse.observation.metadata.usage_object",
"langfuse.observation.usage_details",
}
)
@ -112,14 +110,6 @@ def _expected_generation(route: respx.Route, fixture_name: str, trace_id: str) -
expected_path: Final = Path(__file__).with_name("langfuse_expected_request_body") / fixture_name
expected: Final = TypeAdapter(dict[str, object]).validate_python(json.loads(expected_path.read_text()))
actual: Final = _comparable(generations[0])
attributes: Final = TypeAdapter(dict[str, object]).validate_python(generations[0]["attributes"])
cost_details: Final = TypeAdapter(dict[str, float]).validate_python(attributes["langfuse.observation.cost_details"])
usage_details: Final = TypeAdapter(dict[str, int]).validate_python(attributes["langfuse.observation.usage_details"])
assert cost_details["total"] >= 0
assert usage_details["input"] >= 0
assert usage_details["output"] >= 0
assert usage_details["total"] >= usage_details["input"]
assert usage_details["total"] >= usage_details["output"]
expected_comparable: Final = _comparable(expected)
assert actual == expected_comparable, json.dumps(
{"actual": actual, "expected": expected_comparable}, indent=2, sort_keys=True
@ -449,6 +439,13 @@ async def test_langfuse_logging_with_router(langfuse_export: tuple[respx.Route,
await router.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hello!"}],
mock_response=litellm.ModelResponse(
choices=[],
usage=litellm.Usage(prompt_tokens=10, completion_tokens=10, total_tokens=20),
model="gpt-3.5-turbo",
object="chat.completion",
created=1723081200,
).model_dump(),
metadata={"trace_id": trace_id},
)
await _flush_langfuse()

View file

@ -53,9 +53,8 @@ from litellm.types.services import ServiceLoggerPayload, ServiceTypes
from collections.abc import AsyncIterator
from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.types.utils import ModelResponse
from litellm.types.utils import ChatCompletionMessageToolCall, Function
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
from tests.logging_callback_tests.base_test import BaseLoggingCallbackTest
exporter = InMemorySpanExporter()
@ -7140,9 +7139,20 @@ async def test_arize_phoenix_creates_nested_spans_on_dedicated_provider():
@pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown")
class TestOpentelemetryUnitTests(BaseLoggingCallbackTest):
def test_parallel_tool_calls(self, mock_response_obj: ModelResponse):
tool_calls = mock_response_obj.choices[0].message.tool_calls
class TestOpentelemetryUnitTests:
def test_parallel_tool_calls(self):
tool_calls: Final = [
ChatCompletionMessageToolCall(
function=Function(arguments='{"city": "New York"}', name="get_weather"),
id="call_weather",
type="function",
),
ChatCompletionMessageToolCall(
function=Function(arguments='{"city": "New York"}', name="get_news"),
id="call_news",
type="function",
),
]
from litellm.integrations.opentelemetry import OpenTelemetry
from litellm.proxy._types import SpanAttributes

View file

@ -1,5 +1,4 @@
import asyncio
import os
from pathlib import Path
from unittest import mock
@ -102,16 +101,9 @@ def test_azure_image_generation_route(client_no_auth):
def test_azure_image_edit_route(client_no_auth):
litellm.turn_on_debug()
client, _, mock_aimage_edit = client_no_auth
image_path = os.path.join(
os.path.dirname(__file__),
"../../../image_gen_tests/test_image.png",
)
with open(image_path, "rb") as f:
files = {"image": ("test_image.png", f, "image/png")}
data = {"prompt": "A cute baby sea otter"}
response = client.post(
"/openai/deployments/dall-e-3/images/edits", files=files, data=data
)
files = {"image": ("test_image.png", b"\x89PNG\r\n\x1a\nimage-bytes", "image/png")}
data = {"prompt": "A cute baby sea otter"}
response = client.post("/openai/deployments/dall-e-3/images/edits", files=files, data=data)
mock_aimage_edit.assert_called_once()
called_kwargs = mock_aimage_edit.call_args.kwargs

View file

@ -4216,6 +4216,53 @@ class TestAddMCPServerAtomicity:
assert result.server_id == "created-1"
assert result.rpm == 5
@pytest.mark.asyncio
async def test_create_response_hides_stored_credentials(self):
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
add_mcp_server,
)
payload: Final = NewMCPServerRequest(
alias="echo",
url="https://echo.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.api_key,
credentials={"auth_value": "secret"},
)
admin: Final = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user")
created_server: Final = generate_mock_mcp_server_db_record(
server_id="created-1", alias="echo", auth_type="api_key"
).model_copy(update={"credentials": {"auth_value": "secret"}})
mock_manager: Final = MagicMock()
mock_manager.add_server = AsyncMock()
mock_manager.reload_servers_from_database = AsyncMock()
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=MagicMock(),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload",
MagicMock(),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server_if_identifier_free",
AsyncMock(return_value=created_server),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
mock_manager,
),
):
result: Final = await add_mcp_server(payload=payload, user_api_key_dict=admin)
assert result.server_id == "created-1"
assert result.credentials is None
assert created_server.credentials == {"auth_value": "secret"}
mock_manager.add_server.assert_awaited_once_with(created_server)
@pytest.mark.asyncio
async def test_create_500s_and_skips_registry_when_db_write_fails(self):
from litellm.proxy.management_endpoints.mcp_management_endpoints import (