test_card_resolver_fallback_from_new_to_old_path

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