From 867f097a00b6a53e603c1040682c7fa3e7ae102f Mon Sep 17 00:00:00 2001
From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Date: Wed, 30 Sep 2026 18:57:18 -0700
Subject: [PATCH] test(agents): use real policy data in header forwarding
 fixtures

---
 .../agent_endpoints/test_agent_headers.py     | 27 ++++++++++---------
 1 file changed, 15 insertions(+), 12 deletions(-)

diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py
index e894f4ad69a..4d6077b3472 100644
--- a/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py
+++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py
@@ -14,23 +14,26 @@ from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
+from litellm.types.agents import AgentResponse
+
 # ---------------------------------------------------------------------------
-# Helper: build a minimal mock agent
+# Helper: build a minimal agent
 # ---------------------------------------------------------------------------
 
 
 def _make_mock_agent(
-    static_headers=None,
-    extra_headers=None,
-    url="http://backend-agent:10001",
-):
-    mock_agent = MagicMock()
-    mock_agent.agent_id = "agent-123"
-    mock_agent.agent_card_params = {"url": url, "name": "Test Agent"}
-    mock_agent.litellm_params = {}
-    mock_agent.static_headers = static_headers or {}
-    mock_agent.extra_headers = extra_headers or []
-    return mock_agent
+    static_headers: dict[str, str] | None = None,
+    extra_headers: list[str] | None = None,
+    url: str = "http://backend-agent:10001",
+) -> AgentResponse:
+    return AgentResponse(
+        agent_id="agent-123",
+        agent_name="Test Agent",
+        agent_card_params={"url": url, "name": "Test Agent"},
+        litellm_params={},
+        static_headers=static_headers or {},
+        extra_headers=extra_headers or [],
+    )
 
 
 def _make_mock_request(extra_headers=None, method="message/send"):