scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

test(a2a): make message method test helpers immutable

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-08 22:43:22 +00:00
parent 71adc7542f
commit ec187dd7cf
1 changed files with 38 additions and 50 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

88
tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py
View file

@ -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"