mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
140 lines
4.9 KiB
Python
140 lines
4.9 KiB
Python
"""
|
|
Mock tests for LiteLLMA2ACardResolver.
|
|
|
|
Tests that the card resolver tries both old and new well-known paths.
|
|
"""
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from litellm.a2a_protocol.card_resolver import (
|
|
LiteLLMA2ACardResolver,
|
|
fix_agent_card_url,
|
|
is_localhost_or_internal_url,
|
|
normalize_agent_card_interfaces,
|
|
set_agent_card_url,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_card_resolver_fallback_from_new_to_old_path():
|
|
"""
|
|
Test that the card resolver tries the new path (/.well-known/agent-card.json) first,
|
|
and falls back to the old path (/.well-known/agent.json) if the new path fails.
|
|
"""
|
|
# Mock the AgentCard
|
|
mock_agent_card = MagicMock()
|
|
mock_agent_card.name = "Test Agent"
|
|
mock_agent_card.description = "A test agent"
|
|
|
|
# Track which paths were called
|
|
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
|
|
):
|
|
paths_called.append(relative_card_path)
|
|
if relative_card_path == "/.well-known/agent-card.json":
|
|
# First call (new path) fails
|
|
raise Exception("404 Not Found")
|
|
else:
|
|
# Second call (old path) succeeds
|
|
return mock_agent_card
|
|
|
|
# Create a mock httpx client
|
|
mock_httpx_client = MagicMock()
|
|
|
|
# Patch the parent class's get_agent_card method
|
|
# We need to patch the actual parent class method that super() calls
|
|
with patch.object(
|
|
LiteLLMA2ACardResolver.__bases__[0],
|
|
"get_agent_card",
|
|
mock_parent_get_agent_card,
|
|
):
|
|
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
|
|
assert len(paths_called) == 2
|
|
assert paths_called[0] == "/.well-known/agent-card.json" # New path tried first
|
|
assert paths_called[1] == "/.well-known/agent.json" # Old path tried second
|
|
|
|
# Verify the result
|
|
assert result == mock_agent_card
|
|
assert result.name == "Test Agent"
|
|
|
|
|
|
def test_is_localhost_or_internal_url():
|
|
"""Test that localhost/internal URLs are correctly detected."""
|
|
# Should return True for localhost variants
|
|
assert is_localhost_or_internal_url("http://localhost:8000/") is True
|
|
assert is_localhost_or_internal_url("http://0.0.0.0:8001/") is True
|
|
|
|
# Should return False for public URLs
|
|
assert is_localhost_or_internal_url("https://my-agent.example.com/") is False
|
|
assert is_localhost_or_internal_url(None) is False
|
|
|
|
|
|
def test_fix_agent_card_url_replaces_localhost():
|
|
"""Test that fix_agent_card_url replaces localhost URLs with base_url."""
|
|
# Create mock agent card with localhost URL
|
|
mock_card = MagicMock()
|
|
mock_card.url = "http://0.0.0.0:8001/"
|
|
|
|
# Fix the URL
|
|
result = fix_agent_card_url(mock_card, "https://my-public-agent.example.com")
|
|
|
|
# Verify localhost URL was replaced with base_url
|
|
assert result.url == "https://my-public-agent.example.com/"
|
|
|
|
|
|
def test_set_agent_card_url_updates_top_level_and_supported_interface():
|
|
card = SimpleNamespace(
|
|
url="http://localhost:10001/",
|
|
supported_interfaces=[SimpleNamespace(url="http://0.0.0.0:10001/")],
|
|
)
|
|
|
|
set_agent_card_url(card, "https://my-public-agent.example.com")
|
|
|
|
assert card.url == "https://my-public-agent.example.com/"
|
|
assert card.supported_interfaces[0].url == "https://my-public-agent.example.com/"
|
|
|
|
|
|
def test_fix_agent_card_url_updates_interface_when_top_level_is_localhost():
|
|
card = SimpleNamespace(
|
|
url="http://localhost:10001/",
|
|
supported_interfaces=[SimpleNamespace(url="http://0.0.0.0:10001/")],
|
|
)
|
|
|
|
result = fix_agent_card_url(card, "https://my-public-agent.example.com")
|
|
|
|
assert result.url == "https://my-public-agent.example.com/"
|
|
assert result.supported_interfaces[0].url == "https://my-public-agent.example.com/"
|
|
|
|
|
|
def test_normalize_agent_card_interfaces_downgrades_miscased_interfaces_to_the_0_3_dialect():
|
|
pb2 = pytest.importorskip("a2a.types.a2a_pb2")
|
|
|
|
card = pb2.AgentCard(
|
|
name="langgraph",
|
|
supported_interfaces=[
|
|
pb2.AgentInterface(url="http://a/", protocol_binding="jsonrpc", protocol_version="1.0"),
|
|
pb2.AgentInterface(url="http://b/", protocol_binding="JSONRPC", protocol_version="1.0"),
|
|
pb2.AgentInterface(url="http://c/", protocol_binding="websocket", protocol_version="1.0"),
|
|
],
|
|
)
|
|
|
|
normalized = normalize_agent_card_interfaces(card)
|
|
|
|
assert [(i.protocol_binding, i.protocol_version) for i in normalized.supported_interfaces] == [
|
|
("JSONRPC", "0.3"),
|
|
("JSONRPC", "1.0"),
|
|
("websocket", "1.0"),
|
|
]
|
|
assert card.supported_interfaces[0].protocol_binding == "jsonrpc"
|
|
assert card.supported_interfaces[0].protocol_version == "1.0"
|