From fa74073ea9acc3bc2dedd82dc43c3209251a152c Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 10 Oct 2026 03:38:38 -0700 Subject: [PATCH] test: restore collected no-match cases, real langfuse and gcs payload checks, and gate the aws secret lane --- .../llm_translation/test_ocr_matrix_e2e.py | 2 + .../test_aws_secret_manager_e2e.py | 11 ++++- .../management/test_sso_sign_in.py | 9 +++- tests/integration/mcp/test_mcp_credentials.py | 23 +++++++++ .../integrations/gcs_pubsub/test_pub_sub.py | 48 ++++++++----------- .../integrations/langfuse/test_langfuse.py | 18 +++++-- .../semantic_guard/test_semantic_guard.py | 46 +++++++++--------- .../test_mcp_management_endpoints.py | 47 ------------------ 8 files changed, 101 insertions(+), 103 deletions(-) diff --git a/tests/e2e/llm_translation/test_ocr_matrix_e2e.py b/tests/e2e/llm_translation/test_ocr_matrix_e2e.py index 91ec7649953..0746cc5de4d 100644 --- a/tests/e2e/llm_translation/test_ocr_matrix_e2e.py +++ b/tests/e2e/llm_translation/test_ocr_matrix_e2e.py @@ -13,6 +13,7 @@ from types import MappingProxyType from typing import Final, Literal, Protocol, cast import pytest +from e2e_metadata import Domain, Route, Subject, meta from pydantic import TypeAdapter import litellm @@ -310,6 +311,7 @@ def _assert_logged(logged: LoggedCall, response: OCRResponse, model: str, logged assert response_cost > 0 +@meta(Subject(domain=Domain.LLM_TRANSLATION, route=Route.OCR)) @pytest.mark.parametrize("case", CASES, ids=[case.id for case in CASES]) async def test_ocr( case: Case, diff --git a/tests/e2e/secret_manager/test_aws_secret_manager_e2e.py b/tests/e2e/secret_manager/test_aws_secret_manager_e2e.py index ce3cdb0ac8c..0c6ca0cacaf 100644 --- a/tests/e2e/secret_manager/test_aws_secret_manager_e2e.py +++ b/tests/e2e/secret_manager/test_aws_secret_manager_e2e.py @@ -1,14 +1,23 @@ import hashlib +import os from typing import Final import pytest +from e2e_config import SECRET_MANAGER_OPT_IN_ENV from e2e_metadata import Domain, Subject, meta from pydantic import TypeAdapter from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2 from litellm.secret_managers.main import get_secret -pytestmark = pytest.mark.e2e +pytestmark = [ + pytest.mark.e2e, + pytest.mark.secret_manager, + pytest.mark.skipif( + os.environ.get(SECRET_MANAGER_OPT_IN_ENV, "").strip() != "aws", + reason=f"runs only on the AWS Secrets Manager lane ({SECRET_MANAGER_OPT_IN_ENV}=aws)", + ), +] _AWS_FIXTURE_MASTER_KEY_SHA256: Final = "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b" _STRING_MAP: Final = TypeAdapter(dict[str, str]) diff --git a/tests/integration/management/test_sso_sign_in.py b/tests/integration/management/test_sso_sign_in.py index a441b1fc486..628f1807d55 100644 --- a/tests/integration/management/test_sso_sign_in.py +++ b/tests/integration/management/test_sso_sign_in.py @@ -1,4 +1,5 @@ import base64 +import contextlib import json import uuid from collections.abc import Iterator, Mapping @@ -93,6 +94,11 @@ def _sign_in(gateway: Gateway, subject: str) -> httpx.Response: return browser.get(at_idp.headers["location"]) +def _delete_user(gateway: Gateway, subject: str) -> None: + response: Final = gateway.request("POST", "/user/delete", {"user_ids": [subject]}) + assert response.status_code == 200, response.text + + def _signed_in_user(subject: str) -> list[dict[str, JsonValue]]: return read_rows('SELECT user_email, user_role, metadata FROM "LiteLLM_UserTable" WHERE user_id = %s', (subject,)) @@ -109,7 +115,8 @@ def test_a_first_sso_sign_in_creates_the_user_with_the_default_role( tmp_path: Path, litellm_settings: Mapping[str, JsonValue], expected_role: str ) -> None: subject: Final = f"sso-new-user-{uuid.uuid4().hex[:12]}" - with _sso_gateway(tmp_path, litellm_settings) as gateway: + with _sso_gateway(tmp_path, litellm_settings) as gateway, contextlib.ExitStack() as cleanups: + cleanups.callback(_delete_user, gateway, subject) callback: Final = _sign_in(gateway, subject) assert callback.status_code == 303, f"{callback.status_code} {callback.text}" diff --git a/tests/integration/mcp/test_mcp_credentials.py b/tests/integration/mcp/test_mcp_credentials.py index da952b216d5..208b33b3580 100644 --- a/tests/integration/mcp/test_mcp_credentials.py +++ b/tests/integration/mcp/test_mcp_credentials.py @@ -17,6 +17,7 @@ from integration._support.mcp import ( Outcome, ScriptedTool, call_tool, + forget_mcp, mcp_peer, openapi_peer, register_mcp, @@ -77,6 +78,28 @@ def test_static_credential_reaches_the_peer_in_its_mode_shape_and_is_encrypted_a assert _plaintext_rows(identity, secret) == [], "credential stored in plaintext" +def test_create_response_does_not_echo_the_stored_credential(gateway: Gateway) -> None: + with mcp_peer() as peer, gateway.scenario() as scenario: + secret: Final = "create-secret-" + uuid.uuid4().hex + alias: Final = "cred" + uuid.uuid4().hex[:8] + response: Final = gateway.request( + "POST", + "/v1/mcp/server", + { + "server_name": alias, + "alias": alias, + **peer.registration(), + "auth_type": "bearer_token", + "credentials": {"auth_value": secret}, + }, + ) + scenario.cleanups.callback(forget_mcp, gateway, response.json()["server_id"]) + assert response.status_code == 201, response.text + assert response.json()["alias"] == alias + assert response.json().get("credentials") is None + assert secret not in response.text + + def test_editing_the_credential_rotates_what_the_peer_receives(gateway: Gateway) -> None: with mcp_peer() as peer, gateway.scenario() as scenario: first: Final = "cred-" + uuid.uuid4().hex diff --git a/tests/unit/integrations/gcs_pubsub/test_pub_sub.py b/tests/unit/integrations/gcs_pubsub/test_pub_sub.py index 0000b778626..77fc51a9026 100644 --- a/tests/unit/integrations/gcs_pubsub/test_pub_sub.py +++ b/tests/unit/integrations/gcs_pubsub/test_pub_sub.py @@ -1,3 +1,4 @@ +import asyncio import base64 import datetime import json @@ -16,6 +17,7 @@ import respx import litellm from pydantic import TypeAdapter +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.types.utils import StandardLoggingPayload @@ -200,6 +202,11 @@ def _without_volatile_gcs_payload_fields(value: object, path: str = "") -> objec return value +def _with_decoded_metadata(payload: Mapping[str, object]) -> dict[str, object]: + metadata: Final = payload["metadata"] + return {**payload, "metadata": json.loads(metadata) if isinstance(metadata, str) else metadata} + + @pytest.mark.asyncio async def test_v1_success_event_publishes_the_spend_log_fixture( monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, tmp_path: Path @@ -217,40 +224,25 @@ async def test_v1_success_event_publishes_the_spend_log_fixture( topic_id=topic_id, credentials_path=_authorized_user_credentials(monkeypatch, tmp_path, project_id), ) - standard_logging_payload: Final = cast( - StandardLoggingPayload, - { - "model": "gpt-4o", - "prompt_tokens": 1, - "completion_tokens": 1, - "total_tokens": 2, - "metadata": {}, - "hidden_params": {}, - "model_map_information": None, - }, - ) - response_obj: Final = { - "id": "unit-response", - "model": "gpt-4o", - "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, - "choices": [{"message": {"content": "hi"}}], - } - await logger.async_log_success_event( - kwargs={"model": "gpt-4o", "standard_logging_object": standard_logging_payload}, - response_obj=response_obj, - start_time=datetime.datetime(2025, 1, 1), - end_time=datetime.datetime(2025, 1, 1, 0, 0, 1), + monkeypatch.setattr(litellm, "callbacks", [logger]) + await litellm.acompletion( + model="gpt-4o", + messages=[{"role": "user", "content": "Hello, world!"}], + mock_response="hi", ) + await asyncio.sleep(0) + await GLOBAL_LOGGING_WORKER.flush() await logger.async_send_batch() request_json: Final = TypeAdapter(dict[str, list[dict[str, str]]]).validate_json(route.calls.last.request.content) decoded_payload: Final = TypeAdapter(dict[str, object]).validate_json( base64.b64decode(request_json["messages"][0]["data"]) ) - + expected_payload: Final = TypeAdapter(dict[str, object]).validate_json( + Path(__file__).with_name("spend_logs_payload.json").read_text() + ) assert route.call_count == 1 assert route.calls.last.request.url == url - assert decoded_payload["model"] == "gpt-4o" - assert decoded_payload["prompt_tokens"] == 1 - assert decoded_payload["completion_tokens"] == 1 - assert decoded_payload["total_tokens"] == 2 + assert _without_volatile_gcs_payload_fields(_with_decoded_metadata(decoded_payload)) == ( + _without_volatile_gcs_payload_fields(_with_decoded_metadata(expected_payload)) + ) diff --git a/tests/unit/integrations/langfuse/test_langfuse.py b/tests/unit/integrations/langfuse/test_langfuse.py index 9ef6c9eda85..7a374c1affe 100644 --- a/tests/unit/integrations/langfuse/test_langfuse.py +++ b/tests/unit/integrations/langfuse/test_langfuse.py @@ -294,8 +294,20 @@ async def test_langfuse_logging_with_non_serializable_metadata( ({"a": {"nested_a": 1}, "b": {"nested_b": 2}}, "nested_metadata.json"), ({"a": [1, 2, 3], "b": {4, 5, 6}}, "simple_metadata2.json"), ({"a": (1, 2), "b": frozenset({3, 4}), "c": {"d": [5, 6]}}, "simple_metadata3.json"), - ({"lock": "non-copyable"}, "metadata_with_lock.json"), - ({"func": "callable"}, "metadata_with_function.json"), + ({"lock": threading.Lock()}, "metadata_with_lock.json"), + ({"func": lambda x: x + 1}, "metadata_with_function.json"), + ( + { + "int": 42, + "str": "hello", + "list": [1, 2, 3], + "set": {4, 5}, + "dict": {"nested": "value"}, + "non_copyable": threading.Lock(), + "function": print, + }, + "complex_metadata.json", + ), ({"list": ["list", "not", "a", "dict"]}, "complex_metadata_2.json"), ({}, "empty_metadata.json"), ], @@ -456,7 +468,7 @@ async def test_langfuse_logging_with_router(langfuse_export: tuple[respx.Route, @pytest.mark.asyncio async def test_langfuse_e2e_sync(langfuse_export: tuple[respx.Route, str]) -> None: route, trace_id = langfuse_export - await litellm.acompletion( + litellm.completion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "sync callback route"}], mock_response="Hello from litellm", diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/semantic_guard/test_semantic_guard.py b/tests/unit/proxy/guardrails/guardrail_hooks/semantic_guard/test_semantic_guard.py index 8dec9693726..aa43a5530f3 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/semantic_guard/test_semantic_guard.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/semantic_guard/test_semantic_guard.py @@ -459,6 +459,29 @@ class TestContentFilterPromptInjectionTemplate: ) assert result is None or result["texts"][0] == sentence + NO_MATCH_CASES = [ + ("summarize our Q3 financial results", "no match: normal business query"), + ("help me draft an email to a customer", "no match: normal business query"), + ("what is the capital of the UAE?", "no match: general knowledge"), + ("write a Python function to sort a list", "no match: coding help"), + ("how does a firewall work?", "no match: security education"), + ] + + @pytest.mark.parametrize( + "sentence,reason", + NO_MATCH_CASES, + ids=[f"no_match_{i}" for i in range(len(NO_MATCH_CASES))], + ) + @pytest.mark.asyncio + async def test_no_match_allowed(self, content_filter_guardrail, sentence, reason): + request_data = {"messages": [{"role": "user", "content": sentence}]} + result = await content_filter_guardrail.apply_guardrail( + inputs={"texts": [sentence]}, + request_data=request_data, + input_type="request", + ) + assert result is None or result["texts"][0] == sentence + def test_build_routes_empty() -> None: routes: Final = SemanticGuardRouteLoader.build_routes( @@ -538,26 +561,3 @@ def test_empty_routes_raises() -> None: route_templates=None, custom_routes=None, ) - - NO_MATCH_CASES = [ - ("summarize our Q3 financial results", "no match: normal business query"), - ("help me draft an email to a customer", "no match: normal business query"), - ("what is the capital of the UAE?", "no match: general knowledge"), - ("write a Python function to sort a list", "no match: coding help"), - ("how does a firewall work?", "no match: security education"), - ] - - @pytest.mark.parametrize( - "sentence,reason", - NO_MATCH_CASES, - ids=[f"no_match_{i}" for i in range(len(NO_MATCH_CASES))], - ) - @pytest.mark.asyncio - async def test_no_match_allowed(self, content_filter_guardrail, sentence, reason): - request_data = {"messages": [{"role": "user", "content": sentence}]} - result = await content_filter_guardrail.apply_guardrail( - inputs={"texts": [sentence]}, - request_data=request_data, - input_type="request", - ) - assert result is None or result["texts"][0] == sentence diff --git a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py index 1fee3d0ab4a..21ff5534e03 100644 --- a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -4216,53 +4216,6 @@ 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 (