diff --git a/tests/test_litellm/a2a_protocol/test_card_resolver.py b/tests/test_litellm/a2a_protocol/test_card_resolver.py index 84709bab139..1bdab50860c 100644 --- a/tests/test_litellm/a2a_protocol/test_card_resolver.py +++ b/tests/test_litellm/a2a_protocol/test_card_resolver.py @@ -4,12 +4,12 @@ Mock tests for LiteLLMA2ACardResolver. Tests that the card resolver tries both old and new well-known paths. """ -import sys from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm.a2a_protocol.card_resolver import ( + LiteLLMA2ACardResolver, fix_agent_card_url, is_localhost_or_internal_url, ) @@ -29,39 +29,31 @@ async def test_card_resolver_fallback_from_new_to_old_path(): # Track which paths were called paths_called = [] - # Create a mock base class - class MockA2ACardResolver: - def __init__(self, base_url): - self.base_url = base_url - - async def 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 mock A2A module - mock_a2a_module = MagicMock() - mock_a2a_client = MagicMock() - mock_a2a_constants = MagicMock() - mock_a2a_constants.AGENT_CARD_WELL_KNOWN_PATH = "/.well-known/agent-card.json" - mock_a2a_constants.PREV_AGENT_CARD_WELL_KNOWN_PATH = "/.well-known/agent.json" - - with patch.dict( - sys.modules, - { - "a2a": mock_a2a_module, - "a2a.client": MagicMock(A2ACardResolver=MockA2ACardResolver), - "a2a.utils.constants": mock_a2a_constants, - }, + # 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 ): - # Import after patching - from litellm.a2a_protocol.card_resolver import LiteLLMA2ACardResolver + 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 - resolver = LiteLLMA2ACardResolver(base_url="http://test-agent:8000") + # 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