From b1695e909070562dffb6d01655676a2fa0b8743a Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 04:33:55 -0700
Subject: [PATCH] test(agents): type the public-agent regression tests fully

---
 .../proxy/agent_endpoints/test_endpoints.py   | 25 ++++++++++---------
 1 file changed, 13 insertions(+), 12 deletions(-)

diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py
index 7c1db851a48..51797cbb308 100644
--- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py
+++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py
@@ -1,3 +1,4 @@
+from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -1069,12 +1070,12 @@ class _DbBackedProxyConfig:
     `get_config()`, which is what re-assigns the `litellm.public_*` globals in production."""
 
     def __init__(self) -> None:
-        self.stored_litellm_settings: dict = {}
+        self.stored_litellm_settings: dict[str, object] = {}
 
-    async def get_config(self) -> dict:
+    async def get_config(self) -> dict[str, dict[str, object]]:
         from litellm.proxy.proxy_server import ProxyConfig
 
-        config: dict = {"litellm_settings": {}}
+        config: Final[dict[str, dict[str, object]]] = {"litellm_settings": {}}
         if not self.stored_litellm_settings:
             return config
         return ProxyConfig()._update_config_fields(
@@ -1083,17 +1084,17 @@ class _DbBackedProxyConfig:
             db_param_value=dict(self.stored_litellm_settings),
         )
 
-    async def save_config(self, new_config: dict) -> None:
+    async def save_config(self, new_config: dict[str, dict[str, object]]) -> None:
         self.stored_litellm_settings = dict(new_config.get("litellm_settings") or {})
 
 
-def test_make_agent_public_twice_keeps_both_agents_public(monkeypatch):
+def test_make_agent_public_twice_keeps_both_agents_public(monkeypatch: pytest.MonkeyPatch) -> None:
     """A second /make_public call must not drop the agent published by the first one."""
     import litellm
     from litellm.proxy.agent_endpoints import agent_registry as agent_registry_module
     from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry
 
-    registry = AgentRegistry()
+    registry: Final = AgentRegistry()
     registry.register_agent(_sample_agent_response(agent_id="agent-1", agent_name="Agent One"))
     registry.register_agent(_sample_agent_response(agent_id="agent-2", agent_name="Agent Two"))
 
@@ -1102,8 +1103,8 @@ def test_make_agent_public_twice_keeps_both_agents_public(monkeypatch):
     monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock())
     monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", _DbBackedProxyConfig())
 
-    first = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"})
-    second = client.post("/v1/agents/agent-2/make_public", headers={"Authorization": "Bearer test-key"})
+    first: Final = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"})
+    second: Final = client.post("/v1/agents/agent-2/make_public", headers={"Authorization": "Bearer test-key"})
 
     assert first.status_code == 200
     assert second.status_code == 200
@@ -1111,13 +1112,13 @@ def test_make_agent_public_twice_keeps_both_agents_public(monkeypatch):
     assert [agent.agent_id for agent in registry.get_public_agent_list()] == ["agent-1", "agent-2"]
 
 
-def test_make_agent_public_rejects_an_already_public_agent(monkeypatch):
+def test_make_agent_public_rejects_an_already_public_agent(monkeypatch: pytest.MonkeyPatch) -> None:
     """The duplicate guard must still fire when the published list comes back from the DB."""
     import litellm
     from litellm.proxy.agent_endpoints import agent_registry as agent_registry_module
     from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry
 
-    registry = AgentRegistry()
+    registry: Final = AgentRegistry()
     registry.register_agent(_sample_agent_response(agent_id="agent-1", agent_name="Agent One"))
 
     monkeypatch.setattr(agent_registry_module, "global_agent_registry", registry)
@@ -1125,8 +1126,8 @@ def test_make_agent_public_rejects_an_already_public_agent(monkeypatch):
     monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock())
     monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", _DbBackedProxyConfig())
 
-    first = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"})
-    duplicate = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"})
+    first: Final = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"})
+    duplicate: Final = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"})
 
     assert first.status_code == 200
     assert duplicate.status_code == 400