From ec187dd7cfcb66e5fbbb8e418cfeac144467756e Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 8 Sep 2026 22:43:22 +0000 Subject: [PATCH] test(a2a): make message method test helpers immutable Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../agent_endpoints/test_a2a_endpoints.py | 88 ++++++++----------- 1 file changed, 38 insertions(+), 50 deletions(-) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py index 65174a9e2d2..e0476361074 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py @@ -9,7 +9,8 @@ import socket import sys from collections.abc import Awaitable, Callable, Mapping from contextlib import AbstractContextManager, ExitStack -from typing import TypedDict +from dataclasses import dataclass +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -19,7 +20,8 @@ from litellm.proxy._types import UserAPIKeyAuth AddLiteLLMData = Callable[..., Awaitable[dict[str, object]]] -class CapturedAgentCall(TypedDict): +@dataclass(frozen=True, slots=True) +class CapturedAgentCall: request_id: object agent_extra_headers: dict[str, str] | None @@ -413,14 +415,11 @@ def _base_patches( async def _add_proxy_data(data: dict[str, object], **kwargs: object) -> dict[str, object]: - data["proxy_server_request"] = { - "url": "http://localhost:4000", - "method": "POST", - "headers": {}, - "body": {}, + return { + **data, + "proxy_server_request": {"url": "http://localhost:4000", "method": "POST", "headers": {}, "body": {}}, + "metadata": data.get("metadata", {}), } - data.setdefault("metadata", {}) - return data _HELLO_MESSAGE_PARAMS = { @@ -448,14 +447,8 @@ async def _invoke_message_method( def __init__(self, **kwargs: object) -> None: self.__dict__.update(kwargs) - captured: CapturedAgentCall = {"request_id": None, "agent_extra_headers": None} - - async def capture_asend_message( - request: SendMessageRequest, agent_extra_headers: dict[str, str] | None = None, **kwargs: object - ) -> MagicMock: - captured["request_id"] = request.__dict__["id"] - captured["agent_extra_headers"] = agent_extra_headers - response = MagicMock() + async def fake_asend_message(request: SendMessageRequest, **kwargs: object) -> MagicMock: + response: Final = MagicMock() response.model_dump.return_value = { "jsonrpc": "2.0", "id": request.__dict__["id"], @@ -463,40 +456,25 @@ async def _invoke_message_method( } return response - async def capture_stream_message( - request_id: object, agent_extra_headers: dict[str, str] | None = None, **kwargs: object - ) -> JSONResponse: - captured["request_id"] = request_id - captured["agent_extra_headers"] = agent_extra_headers + async def fake_stream_message(request_id: object, **kwargs: object) -> JSONResponse: return JSONResponse({"jsonrpc": "2.0", "id": request_id}) - mock_a2a_types = MagicMock() + mock_a2a_types: Final = MagicMock() mock_a2a_types.MessageSendParams = MessageSendParams mock_a2a_types.SendMessageRequest = SendMessageRequest + is_send: Final = method == "message/send" + downstream: Final = AsyncMock(side_effect=fake_asend_message if is_send else fake_stream_message) with ExitStack() as stack: for p in _base_patches(_make_agent_mock(), add_litellm_data): stack.enter_context(p) stack.enter_context(patch("litellm.a2a_protocol.main.A2A_SDK_AVAILABLE", True)) - if method == "message/send": - stack.enter_context( - patch.dict( - sys.modules, - {"a2a": MagicMock(), "a2a.types": mock_a2a_types}, - ) - ) - stack.enter_context( - patch( - "litellm.a2a_protocol.asend_message", - new=AsyncMock(side_effect=capture_asend_message), - ) - ) + if is_send: + stack.enter_context(patch.dict(sys.modules, {"a2a": MagicMock(), "a2a.types": mock_a2a_types})) + stack.enter_context(patch("litellm.a2a_protocol.asend_message", new=downstream)) else: stack.enter_context( - patch( - "litellm.proxy.agent_endpoints.a2a_endpoints._handle_stream_message", - new=AsyncMock(side_effect=capture_stream_message), - ) + patch("litellm.proxy.agent_endpoints.a2a_endpoints._handle_stream_message", new=downstream) ) from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a @@ -508,7 +486,9 @@ async def _invoke_message_method( user_api_key_dict=user_api_key_dict, ) - return captured + kwargs: Final = downstream.call_args.kwargs + request_id: Final = kwargs["request"].__dict__["id"] if is_send else kwargs["request_id"] + return CapturedAgentCall(request_id=request_id, agent_extra_headers=kwargs.get("agent_extra_headers")) @pytest.mark.asyncio @@ -519,7 +499,7 @@ async def test_message_methods_preserve_numeric_zero_request_id(method: str): captured = await _invoke_message_method(method, mock_request, user_api_key_dict) - assert captured["request_id"] == 0 + assert captured.request_id == 0 @pytest.mark.asyncio @@ -530,7 +510,7 @@ async def test_message_methods_forward_caller_identity_headers(method: str): captured = await _invoke_message_method(method, mock_request, user_api_key_dict) - forwarded_headers = captured["agent_extra_headers"] or {} + forwarded_headers = captured.agent_extra_headers or {} assert forwarded_headers.get("X-LiteLLM-User-Id") == "user-abc" assert forwarded_headers.get("X-LiteLLM-Team-Id") == "team-xyz" @@ -547,7 +527,7 @@ async def test_message_methods_caller_identity_headers_cannot_be_spoofed(method: captured = await _invoke_message_method(method, mock_request, user_api_key_dict) - forwarded_headers = captured["agent_extra_headers"] or {} + forwarded_headers = captured.agent_extra_headers or {} assert ( forwarded_headers.get("X-LiteLLM-User-Id") == "real-user" ), "authenticated user id must not be overridden by forwarded client headers" @@ -559,19 +539,27 @@ async def test_message_methods_caller_identity_headers_cannot_be_spoofed(method: @pytest.mark.asyncio @pytest.mark.parametrize("method", ["message/send", "message/stream"]) async def test_message_methods_forward_key_bound_identity_not_pre_call_rewrite(method: str): - mock_request = _make_request_mock(method, _HELLO_MESSAGE_PARAMS) - user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="key-user", team_id="key-team") + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup - async def rewrite_user_id_like_header_mapping(data: dict[str, object], **kwargs: object) -> dict[str, object]: - user_api_key_dict.user_id = "header-mapped-user" + mock_request = _make_request_mock(method, _HELLO_MESSAGE_PARAMS) + mock_request.headers = {"X-OpenWebUI-User-Id": "header-mapped-user"} + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="key-user", team_id="key-team") + general_settings: Final = { + "user_header_mappings": [{"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}] + } + + async def apply_user_header_mapping(data: dict[str, object], **kwargs: object) -> dict[str, object]: + LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( + general_settings, user_api_key_dict, dict(mock_request.headers) + ) return await _add_proxy_data(data, **kwargs) captured = await _invoke_message_method( - method, mock_request, user_api_key_dict, add_litellm_data=rewrite_user_id_like_header_mapping + method, mock_request, user_api_key_dict, add_litellm_data=apply_user_header_mapping ) assert user_api_key_dict.user_id == "header-mapped-user", "precondition: pre-call rewrite ran" - forwarded_headers = captured["agent_extra_headers"] or {} + forwarded_headers = captured.agent_extra_headers or {} assert forwarded_headers.get("X-LiteLLM-User-Id") == "key-user" assert forwarded_headers.get("X-LiteLLM-Team-Id") == "key-team"