test(agents): type the public-agent regression tests fully

This commit is contained in:
mateo-berri 2026-09-03 04:33:55 -07:00
parent 9d862a6583
commit b1695e9090

View file

@ -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