scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-25 01:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 131

Merge pull request #42160 from BerriAI/litellm_migrate_tests_p17

Browse source 
test(a2a): migrate a2a_protocol legacy tests to tests/unit (wave 3 phase 17)
This commit is contained in:
yuneng-jiang 2026-09-21 09:19:29 -07:00 • committed by GitHub
commit 27797923a3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 34 additions and 203 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

0
tests/unit/a2a_protocol/providers/bedrock_agentcore/__init__.py Normal file
View file

133
tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py → tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py
View file

@ -10,12 +10,11 @@ Verifies that:
"""
import json
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import respx
from unittest.mock import AsyncMock, MagicMock, patch
SAMPLE_ARN = "arn:aws:bedrock-agentcore:us-west-2:123456789:runtime/my_agent"
SAMPLE_MODEL = f"bedrock/agentcore/{SAMPLE_ARN}"
@ -42,13 +41,11 @@ class TestTransformation:
BedrockAgentCoreA2ATransformation,
)
url, headers, body = (
BedrockAgentCoreA2ATransformation.get_url_and_signed_request(
request_id="req-001",
params=SAMPLE_PARAMS,
litellm_params=SAMPLE_LITELLM_PARAMS,
method="message/send",
)
url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request(
request_id="req-001",
params=SAMPLE_PARAMS,
litellm_params=SAMPLE_LITELLM_PARAMS,
method="message/send",
)
body_dict = json.loads(body)
assert body_dict["jsonrpc"] == "2.0"
@ -201,10 +198,7 @@ class TestTransformation:
# Runtime user id is the value set from litellm_params, NOT the spoof.
assert normalized["x-amzn-bedrock-agentcore-runtime-user-id"] == "legit-user"
# Session id is the auto-generated one, not the spoofed value.
assert (
normalized["x-amzn-bedrock-agentcore-runtime-session-id"]
!= "spoofed-session"
)
assert normalized["x-amzn-bedrock-agentcore-runtime-session-id"] != "spoofed-session"
# Authorization is the JWT bearer set by the signer, not the spoof.
assert normalized["authorization"] == "Bearer test-jwt-token"
# Host / x-amz-* must not have been carried over from the client.
@ -259,43 +253,6 @@ class TestTransformation:
# Non-reserved header still makes it into the signed dict.
assert captured.get("x-mcp-token") == "mcp-abc"
def test_sigv4_auth_when_no_api_key(self):
"""When no api_key, falls through to SigV4 signing."""
from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import (
BedrockAgentCoreA2ATransformation,
)
litellm_params_no_key = {
"model": SAMPLE_MODEL,
"custom_llm_provider": "bedrock",
"aws_access_key_id": "AKIAIOSFODNN7EXAMPLE",
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
"aws_region_name": "us-west-2",
}
# Mock _sign_request to avoid hitting real botocore credential resolution
fake_sigv4_headers = {
"Authorization": "AWS4-HMAC-SHA256 Credential=AKIA.../bedrock-agentcore/aws4_request",
"Content-Type": "application/json",
"Accept": "application/json, text/event-stream",
}
fake_body = b'{"jsonrpc":"2.0"}'
with patch(
"litellm.llms.bedrock.chat.agentcore.transformation.AmazonAgentCoreConfig._sign_request",
return_value=(fake_sigv4_headers, fake_body),
):
_, headers, _ = (
BedrockAgentCoreA2ATransformation.get_url_and_signed_request(
request_id="req-001",
params=SAMPLE_PARAMS,
litellm_params=litellm_params_no_key,
)
)
# SigV4 produces an Authorization header starting with "AWS4-HMAC-SHA256"
assert "Authorization" in headers
assert headers["Authorization"].startswith("AWS4-HMAC-SHA256")
SESSION_HEADER = "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"
CONTEXT_ID = "conversation-alpha-0001-0000000000000000"
@ -571,39 +528,6 @@ class TestNonStreaming:
sent_headers = mock_client.post.call_args.kwargs["headers"]
assert sent_headers.get("x-mcp-token") == "mcp-abc"
@pytest.mark.asyncio
async def test_a2a_error_response_passthrough(self):
"""JSON-RPC error responses from the agent are returned as-is."""
from litellm.a2a_protocol.providers.bedrock_agentcore.config import (
BedrockAgentCoreA2AConfig,
)
error_response = {
"jsonrpc": "2.0",
"id": "req-001",
"error": {"code": -32600, "message": "Bad request"},
}
mock_response = MagicMock()
mock_response.json.return_value = error_response
mock_response.raise_for_status = MagicMock()
with patch(
"litellm.a2a_protocol.providers.bedrock_agentcore.handler.get_async_httpx_client"
) as mock_get_client:
mock_client = AsyncMock()
mock_client.post = AsyncMock(return_value=mock_response)
mock_get_client.return_value = mock_client
config = BedrockAgentCoreA2AConfig()
result = await config.handle_non_streaming(
request_id="req-001",
params=SAMPLE_PARAMS,
litellm_params=SAMPLE_LITELLM_PARAMS,
)
assert result["error"]["code"] == -32600
assert result["error"]["message"] == "Bad request"
class TestStreaming:
"""Streaming requests must ask AgentCore for a stream, not a single send."""
@ -648,9 +572,7 @@ class TestConfigManager:
A2AProviderConfigManager,
)
config = A2AProviderConfigManager.get_provider_config(
"bedrock", model=SAMPLE_MODEL
)
config = A2AProviderConfigManager.get_provider_config("bedrock", model=SAMPLE_MODEL)
assert config is not None
assert isinstance(config, BedrockAgentCoreA2AConfig)
@ -660,9 +582,7 @@ class TestConfigManager:
A2AProviderConfigManager,
)
config = A2AProviderConfigManager.get_provider_config(
"bedrock", model="bedrock/anthropic.claude-3-sonnet"
)
config = A2AProviderConfigManager.get_provider_config("bedrock", model="bedrock/anthropic.claude-3-sonnet")
assert config is None
def test_unknown_provider_returns_none(self):
@ -676,37 +596,6 @@ class TestConfigManager:
class TestHandlerIntegration:
"""Test handler.py changes — litellm_params passed through, api_base not required."""
@pytest.mark.asyncio
async def test_provider_config_receives_litellm_params(self):
"""Verify handler passes litellm_params to provider config via kwargs."""
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
A2ACompletionBridgeHandler,
)
mock_config = AsyncMock()
mock_config.handle_non_streaming = AsyncMock(
return_value={"jsonrpc": "2.0", "id": "req-001", "result": {}}
)
with patch(
"litellm.a2a_protocol.litellm_completion_bridge.handler.A2AProviderConfigManager.get_provider_config",
return_value=mock_config,
):
await A2ACompletionBridgeHandler.handle_non_streaming(
request_id="req-001",
params=SAMPLE_PARAMS,
litellm_params=SAMPLE_LITELLM_PARAMS,
api_base=None,
)
mock_config.handle_non_streaming.assert_called_once_with(
request_id="req-001",
params=SAMPLE_PARAMS,
api_base=None,
litellm_params=SAMPLE_LITELLM_PARAMS,
agent_extra_headers=None,
)
@pytest.mark.asyncio
async def test_api_base_none_allowed_with_provider_config(self):
"""api_base=None no longer raises when a provider config is registered."""
@ -715,9 +604,7 @@ class TestHandlerIntegration:
)
mock_config = AsyncMock()
mock_config.handle_non_streaming = AsyncMock(
return_value={"jsonrpc": "2.0", "id": "req-001", "result": {}}
)
mock_config.handle_non_streaming = AsyncMock(return_value={"jsonrpc": "2.0", "id": "req-001", "result": {}})
with patch(
"litellm.a2a_protocol.litellm_completion_bridge.handler.A2AProviderConfigManager.get_provider_config",

5
tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py → tests/unit/a2a_protocol/test_a2a_exception_mapping_utils.py
View file

@ -38,9 +38,7 @@ async def test_localhost_retry_reuses_stashed_httpx_client():
patch.object(emu, "A2A_SDK_AVAILABLE", True),
patch.object(emu, "set_agent_card_url") as mock_set_url,
patch.object(emu, "ClientConfig", side_effect=fake_client_config),
patch.object(
emu, "create_client", new=AsyncMock(return_value=new_client)
) as mock_create,
patch.object(emu, "create_client", new=AsyncMock(return_value=new_client)) as mock_create,
):
result = await emu.handle_a2a_localhost_retry(
error=_localhost_error(),
@ -171,6 +169,7 @@ async def test_stream_with_retry_raises_after_localhost_retries_exhausted():
api_base="https://agent.example",
agent_name="test-agent",
)
async def _drain():
async for _chunk in stream:
pytest.fail("expected retry exhaustion to raise before yielding")

23
tests/test_litellm/a2a_protocol/test_a2a_streaming_iterator.py → tests/unit/a2a_protocol/test_a2a_streaming_iterator.py
View file

@ -43,25 +43,6 @@ class RecordingExecutor:
return [fn for fn in self.submits if getattr(fn, "__self__", None) is logging_obj]
@pytest.fixture(autouse=True)
def _isolate_callbacks():
saved = (
litellm.callbacks,
litellm.success_callback,
litellm._async_success_callback,
litellm.failure_callback,
litellm._async_failure_callback,
)
yield
(
litellm.callbacks,
litellm.success_callback,
litellm._async_success_callback,
litellm.failure_callback,
litellm._async_failure_callback,
) = saved
@pytest.mark.asyncio
async def test_custom_logger_only_never_submits_sync_success_handler(monkeypatch):
recording_executor = RecordingExecutor(thread_pool_executor_module.executor)
@ -69,8 +50,8 @@ async def test_custom_logger_only_never_submits_sync_success_handler(monkeypatch
monkeypatch.setattr(a2a_streaming_iterator_module, "executor", recording_executor, raising=False)
recorder = RecordingCustomLogger()
litellm.success_callback = [recorder]
litellm._async_success_callback = [recorder]
monkeypatch.setattr(litellm, "success_callback", [recorder])
monkeypatch.setattr(litellm, "_async_success_callback", [recorder])
logging_obj = LitellmLogging(
model="a2a/test-agent",

8
tests/test_litellm/a2a_protocol/test_card_resolver.py → tests/unit/a2a_protocol/test_card_resolver.py
View file

@ -36,9 +36,7 @@ async def test_card_resolver_fallback_from_new_to_old_path():
paths_called = []
# Create a mock for the parent's get_agent_card method
async def mock_parent_get_agent_card(
self, relative_card_path=None, http_kwargs=None
):
async def mock_parent_get_agent_card(self, relative_card_path=None, http_kwargs=None):
paths_called.append(relative_card_path)
if relative_card_path == "/.well-known/agent-card.json":
# First call (new path) fails
@ -57,9 +55,7 @@ async def test_card_resolver_fallback_from_new_to_old_path():
"get_agent_card",
mock_parent_get_agent_card,
):
resolver = LiteLLMA2ACardResolver(
httpx_client=mock_httpx_client, base_url="http://test-agent:8000"
)
resolver = LiteLLMA2ACardResolver(httpx_client=mock_httpx_client, base_url="http://test-agent:8000")
result = await resolver.get_agent_card()
# Verify both paths were tried in correct order

6
tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py → tests/unit/a2a_protocol/test_completion_bridge_streaming.py
View file

@ -344,11 +344,7 @@ async def test_handle_streaming_keeps_agent_card_path_out_of_the_completion_call
chunk.choices[0].delta.content = "Hello"
yield chunk
with (
patch( # test-quality-ok: the bridge calls litellm.acompletion directly; the sibling tests capture its kwargs through the same seam
"litellm.acompletion", new_callable=AsyncMock
) as mock_acompletion
):
with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion:
mock_acompletion.return_value = mock_streaming_response()
events = [

16
tests/test_litellm/a2a_protocol/test_cost_calculator.py → tests/unit/a2a_protocol/test_cost_calculator.py
View file

@ -122,7 +122,7 @@ class CostLogger(CustomLogger):
@pytest.mark.asyncio
async def test_asend_message_uses_cost_per_query():
async def test_asend_message_uses_cost_per_query(monkeypatch):
"""
Test that asend_message uses cost_per_query param for response_cost.
"""
@ -131,7 +131,7 @@ async def test_asend_message_uses_cost_per_query():
# Setup logger
litellm.logging_callback_manager._reset_all_callbacks()
cost_logger = CostLogger()
litellm.callbacks = [cost_logger]
monkeypatch.setattr(litellm, "callbacks", [cost_logger])
# Mock A2A client
mock_client = MagicMock()
@ -157,7 +157,7 @@ async def test_asend_message_uses_cost_per_query():
@pytest.mark.asyncio
async def test_asend_message_uses_cost_per_query_from_litellm_params_dict():
async def test_asend_message_uses_cost_per_query_from_litellm_params_dict(monkeypatch):
"""
Proxy passes agent pricing as the litellm_params dict param (not top-level
kwargs). Regression for cost_per_query landing at $0 on the native path.
@ -166,7 +166,7 @@ async def test_asend_message_uses_cost_per_query_from_litellm_params_dict():
litellm.logging_callback_manager._reset_all_callbacks()
cost_logger = CostLogger()
litellm.callbacks = [cost_logger]
monkeypatch.setattr(litellm, "callbacks", [cost_logger])
mock_client = MagicMock()
mock_client._litellm_agent_card = MagicMock()
@ -217,7 +217,7 @@ class TokenAndCostLogger(CustomLogger):
@pytest.mark.asyncio
async def test_asend_message_uses_input_output_cost_per_token():
async def test_asend_message_uses_input_output_cost_per_token(monkeypatch):
"""
Test that asend_message calculates cost using input_cost_per_token and output_cost_per_token.
Validates exact cost calculation: cost = (prompt_tokens * input_cost) + (completion_tokens * output_cost)
@ -227,7 +227,7 @@ async def test_asend_message_uses_input_output_cost_per_token():
# Setup logger
litellm.logging_callback_manager._reset_all_callbacks()
token_cost_logger = TokenAndCostLogger()
litellm.callbacks = [token_cost_logger]
monkeypatch.setattr(litellm, "callbacks", [token_cost_logger])
# Mock A2A client
mock_client = MagicMock()
@ -292,7 +292,7 @@ class AgentIdLogger(CustomLogger):
@pytest.mark.asyncio
async def test_asend_message_passes_agent_id_to_callback():
async def test_asend_message_passes_agent_id_to_callback(monkeypatch):
"""
Test that asend_message passes agent_id to callbacks via kwargs.
"""
@ -301,7 +301,7 @@ async def test_asend_message_passes_agent_id_to_callback():
# Setup logger
litellm.logging_callback_manager._reset_all_callbacks()
agent_id_logger = AgentIdLogger()
litellm.callbacks = [agent_id_logger]
monkeypatch.setattr(litellm, "callbacks", [agent_id_logger])
# Mock A2A client
mock_client = MagicMock()

17
tests/test_litellm/a2a_protocol/test_main.py → tests/unit/a2a_protocol/test_main.py
View file

@ -115,9 +115,7 @@ async def test_streaming_trace_id_prefers_logging_trace_id():
captured["extra_headers"] = extra_headers
raise RuntimeError("stop")
with patch.object(
a2a_main, "create_a2a_client", new=AsyncMock(side_effect=_capture)
):
with patch.object(a2a_main, "create_a2a_client", new=AsyncMock(side_effect=_capture)):
with pytest.raises(RuntimeError, match="stop"):
async for _ in a2a_main.asend_message_streaming(
request=request,
@ -229,9 +227,7 @@ _LOWERCASE_BINDING_CARD = {
"defaultInputModes": ["text/plain"],
"defaultOutputModes": ["text/plain"],
"skills": [],
"supportedInterfaces": [
{"url": "http://127.0.0.1:9/", "protocolBinding": "jsonrpc", "protocolVersion": "1.0"}
],
"supportedInterfaces": [{"url": "http://127.0.0.1:9/", "protocolBinding": "jsonrpc", "protocolVersion": "1.0"}],
}
@ -289,11 +285,10 @@ async def _seed_shared_a2a_client(
@pytest.fixture
def isolated_client_cache():
previous = getattr(litellm, "in_memory_llm_clients_cache", None)
litellm.in_memory_llm_clients_cache = LLMClientCache()
yield litellm.in_memory_llm_clients_cache
litellm.in_memory_llm_clients_cache = previous
def isolated_client_cache(monkeypatch):
cache = LLMClientCache()
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", cache)
return cache
def _send_request(request_id):

29
tests/test_litellm/a2a_protocol/test_send_message_response.py → tests/unit/a2a_protocol/test_send_message_response.py
View file

@ -9,9 +9,7 @@ def test_from_dict_backfills_id_on_agent_error_response():
"error": {"code": -32054, "message": "Session not found"},
}
response = LiteLLMSendMessageResponse.from_dict(
agent_error, request_id="r1"
)
response = LiteLLMSendMessageResponse.from_dict(agent_error, request_id="r1")
assert response.id == "r1"
assert response.error == {"code": -32054, "message": "Session not found"}
@ -25,9 +23,7 @@ def test_from_dict_preserves_existing_id():
"error": {"code": -32001, "message": "Task not found"},
}
response = LiteLLMSendMessageResponse.from_dict(
payload, request_id="r1"
)
response = LiteLLMSendMessageResponse.from_dict(payload, request_id="r1")
assert response.id == "upstream-id"
@ -82,9 +78,7 @@ def test_from_dict_accepts_null_id_when_the_error_cannot_be_correlated():
"""JSON-RPC 2.0 section 5 requires ``id`` to be null on an error that cannot be
matched to a request, which is exactly the case where the caller supplied no id
for the backfill to use. Rejecting it turned an agent's error into a proxy 500."""
response = LiteLLMSendMessageResponse.from_dict(
{"jsonrpc": "2.0", "error": {"code": -32054, "message": "x"}}
)
response = LiteLLMSendMessageResponse.from_dict({"jsonrpc": "2.0", "error": {"code": -32054, "message": "x"}})
assert response.id is None
assert response.error == {"code": -32054, "message": "x"}
@ -100,23 +94,6 @@ def test_from_dict_accepts_null_id_echoed_by_upstream():
assert response.id is None
def test_id_accepts_every_member_of_the_json_rpc_union_and_nothing_else():
"""One test pinning the whole ``string | integer | null`` union the spec defines,
so widening the annotation cannot silently become "accept anything"."""
for accepted in ("s1", 42, 0, None):
assert LiteLLMSendMessageResponse(id=accepted).id == accepted
# ``True``/``False`` are in here because bool subclasses int: a non-strict integer
# half would accept them and relay them as 1/0. Direct construction bypasses
# normalization, so the model has to hold this line on its own.
for rejected in (True, False, 1.5, ["a"], {"a": 1}):
try:
LiteLLMSendMessageResponse(id=rejected)
except Exception:
continue
raise AssertionError(f"id={rejected!r} is outside the JSON-RPC union and must be rejected")
def test_boolean_id_is_never_relayed_as_an_integer():
"""``bool`` subclasses ``int``, so widening the annotation to accept integers also
made pydantic coerce a boolean id to 1 or 0. That is worse than rejecting it: an id

0
tests/test_litellm/a2a_protocol/test_utils.py → tests/unit/a2a_protocol/test_utils.py
View file