mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
6a0d03914c
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
516 lines
19 KiB
Python
516 lines
19 KiB
Python
"""
|
||
Unit tests for litellm/proxy/google_endpoints/agents_endpoints.py
|
||
|
||
Focus: verify that list_gemini_agents, get_gemini_agent, delete_gemini_agent,
|
||
and list_gemini_agent_versions correctly forward per-request credentials
|
||
(api_key, api_base, …) supplied via the JSON-encoded litellm_params_template
|
||
query parameter. Flat credential query params (e.g. ?api_key=…) are no
|
||
longer accepted — they would appear in server logs.
|
||
"""
|
||
|
||
import json
|
||
from unittest.mock import AsyncMock, MagicMock, patch
|
||
|
||
import pytest
|
||
from fastapi import Request
|
||
from fastapi.datastructures import Headers, QueryParams
|
||
|
||
|
||
from litellm.proxy.google_endpoints.agents_endpoints import (
|
||
_merge_query_params_into_data,
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _make_request(query_string: str = "") -> MagicMock:
|
||
"""Build a minimal mock Request whose query_params match *query_string*."""
|
||
req = MagicMock(spec=Request)
|
||
req.query_params = QueryParams(query_string)
|
||
req.headers = Headers({})
|
||
return req
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# _merge_query_params_into_data – unit tests for the helper
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestMergeQueryParamsIntoData:
|
||
def test_no_query_params_leaves_data_unchanged(self):
|
||
data = {"custom_llm_provider": "gemini"}
|
||
request = _make_request("")
|
||
result = _merge_query_params_into_data(data, request)
|
||
assert result == {"custom_llm_provider": "gemini"}
|
||
|
||
def test_flat_api_key_is_ignored(self):
|
||
"""Flat credential params must NOT be merged (they leak into server logs)."""
|
||
data = {"custom_llm_provider": "gemini"}
|
||
request = _make_request("api_key=AIzaSyTest123")
|
||
_merge_query_params_into_data(data, request)
|
||
assert "api_key" not in data
|
||
assert data["custom_llm_provider"] == "gemini"
|
||
|
||
def test_flat_params_are_silently_dropped(self):
|
||
"""Flat params (including name injection attempts) are ignored entirely."""
|
||
data = {"name": "my-agent", "custom_llm_provider": "gemini"}
|
||
request = _make_request("name=INJECTED&api_key=AIzaSyTest")
|
||
_merge_query_params_into_data(data, request)
|
||
assert data["name"] == "my-agent"
|
||
assert "api_key" not in data
|
||
|
||
def test_litellm_params_template_json_is_expanded(self):
|
||
template = json.dumps(
|
||
{"api_key": "AIzaFromTemplate", "api_base": "https://example.com"}
|
||
)
|
||
from urllib.parse import quote
|
||
|
||
request = _make_request(f"litellm_params_template={quote(template)}")
|
||
data = {"custom_llm_provider": "gemini"}
|
||
_merge_query_params_into_data(data, request)
|
||
assert data["api_key"] == "AIzaFromTemplate"
|
||
assert data["api_base"] == "https://example.com"
|
||
# The raw template key itself must NOT appear in data
|
||
assert "litellm_params_template" not in data
|
||
|
||
def test_litellm_params_template_does_not_overwrite_existing(self):
|
||
template = json.dumps(
|
||
{"api_key": "FromTemplate", "custom_llm_provider": "openai"}
|
||
)
|
||
from urllib.parse import quote
|
||
|
||
request = _make_request(f"litellm_params_template={quote(template)}")
|
||
data = {"custom_llm_provider": "gemini"}
|
||
_merge_query_params_into_data(data, request)
|
||
# custom_llm_provider was already set; template must not override it
|
||
assert data["custom_llm_provider"] == "gemini"
|
||
assert data["api_key"] == "FromTemplate"
|
||
|
||
def test_invalid_litellm_params_template_json_is_ignored(self):
|
||
request = _make_request("litellm_params_template=NOT_VALID_JSON")
|
||
data = {"custom_llm_provider": "gemini"}
|
||
_merge_query_params_into_data(data, request)
|
||
# Bad JSON is silently skipped; other data stays intact
|
||
assert data == {"custom_llm_provider": "gemini"}
|
||
|
||
def test_template_only_no_flat_params_merged(self):
|
||
"""Only litellm_params_template is expanded; unknown flat params are dropped."""
|
||
template = json.dumps({"api_key": "FromTemplate"})
|
||
from urllib.parse import quote
|
||
|
||
qs = f"litellm_params_template={quote(template)}&vertex_project=my-project"
|
||
request = _make_request(qs)
|
||
data = {"custom_llm_provider": "gemini"}
|
||
_merge_query_params_into_data(data, request)
|
||
assert data["api_key"] == "FromTemplate"
|
||
# flat vertex_project is ignored since it wasn't in litellm_params_template
|
||
assert "vertex_project" not in data
|
||
assert "litellm_params_template" not in data
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Endpoint-level smoke tests: data dict is populated before the processor call
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.fixture
|
||
def mock_srv():
|
||
"""Patch _proxy_server_imports to return lightweight fakes."""
|
||
srv = {
|
||
"general_settings": {},
|
||
"llm_router": MagicMock(),
|
||
"proxy_config": MagicMock(),
|
||
"proxy_logging_obj": MagicMock(),
|
||
"select_data_generator": MagicMock(),
|
||
"user_api_base": None,
|
||
"user_max_tokens": None,
|
||
"user_model": None,
|
||
"user_request_timeout": None,
|
||
"user_temperature": None,
|
||
"version": "0.0.0",
|
||
}
|
||
with patch(
|
||
"litellm.proxy.google_endpoints.agents_endpoints._proxy_server_imports",
|
||
return_value=srv,
|
||
):
|
||
yield srv
|
||
|
||
|
||
@pytest.fixture
|
||
def user_api_key_dict():
|
||
from litellm.proxy._types import UserAPIKeyAuth
|
||
|
||
return UserAPIKeyAuth(api_key="test-key")
|
||
|
||
|
||
def _make_endpoint_request(query_string: str = "") -> MagicMock:
|
||
req = MagicMock(spec=Request)
|
||
req.query_params = QueryParams(query_string)
|
||
req.headers = Headers({})
|
||
req.scope = {}
|
||
|
||
async def _body():
|
||
return b""
|
||
|
||
req.body = _body
|
||
return req
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_list_gemini_agents_passes_api_key_to_processor(
|
||
mock_srv, user_api_key_dict
|
||
):
|
||
from urllib.parse import quote
|
||
|
||
from litellm.proxy.google_endpoints.agents_endpoints import list_gemini_agents
|
||
|
||
template = json.dumps({"api_key": "AIzaListTest"})
|
||
|
||
with patch(
|
||
"litellm.proxy.google_endpoints.agents_endpoints.ProxyBaseLLMRequestProcessing"
|
||
) as MockProcessor:
|
||
instance = MockProcessor.return_value
|
||
instance.base_process_llm_request = AsyncMock(return_value=MagicMock())
|
||
|
||
request = _make_endpoint_request(f"litellm_params_template={quote(template)}")
|
||
await list_gemini_agents(
|
||
request=request,
|
||
fastapi_response=MagicMock(),
|
||
user_api_key_dict=user_api_key_dict,
|
||
)
|
||
|
||
init_data = MockProcessor.call_args[1]["data"]
|
||
assert init_data.get("api_key") == "AIzaListTest"
|
||
assert init_data.get("custom_llm_provider") == "gemini"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_gemini_agent_passes_api_key_to_processor(
|
||
mock_srv, user_api_key_dict
|
||
):
|
||
from urllib.parse import quote
|
||
|
||
from litellm.proxy.google_endpoints.agents_endpoints import get_gemini_agent
|
||
|
||
template = json.dumps({"api_key": "AIzaGetTest"})
|
||
|
||
with patch(
|
||
"litellm.proxy.google_endpoints.agents_endpoints.ProxyBaseLLMRequestProcessing"
|
||
) as MockProcessor:
|
||
instance = MockProcessor.return_value
|
||
instance.base_process_llm_request = AsyncMock(return_value=MagicMock())
|
||
|
||
request = _make_endpoint_request(f"litellm_params_template={quote(template)}")
|
||
await get_gemini_agent(
|
||
request=request,
|
||
name="my-agent",
|
||
fastapi_response=MagicMock(),
|
||
user_api_key_dict=user_api_key_dict,
|
||
)
|
||
|
||
init_data = MockProcessor.call_args[1]["data"]
|
||
assert init_data.get("api_key") == "AIzaGetTest"
|
||
assert init_data.get("name") == "my-agent"
|
||
assert init_data.get("custom_llm_provider") == "gemini"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_delete_gemini_agent_passes_api_key_to_processor(
|
||
mock_srv, user_api_key_dict
|
||
):
|
||
from urllib.parse import quote
|
||
|
||
from litellm.proxy.google_endpoints.agents_endpoints import delete_gemini_agent
|
||
|
||
template = json.dumps({"api_key": "AIzaDeleteTest"})
|
||
|
||
with patch(
|
||
"litellm.proxy.google_endpoints.agents_endpoints.ProxyBaseLLMRequestProcessing"
|
||
) as MockProcessor:
|
||
instance = MockProcessor.return_value
|
||
instance.base_process_llm_request = AsyncMock(return_value=MagicMock())
|
||
|
||
request = _make_endpoint_request(f"litellm_params_template={quote(template)}")
|
||
await delete_gemini_agent(
|
||
request=request,
|
||
name="my-agent",
|
||
fastapi_response=MagicMock(),
|
||
user_api_key_dict=user_api_key_dict,
|
||
)
|
||
|
||
init_data = MockProcessor.call_args[1]["data"]
|
||
assert init_data.get("api_key") == "AIzaDeleteTest"
|
||
assert init_data.get("name") == "my-agent"
|
||
assert init_data.get("custom_llm_provider") == "gemini"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_list_gemini_agent_versions_passes_api_key_to_processor(
|
||
mock_srv, user_api_key_dict
|
||
):
|
||
from urllib.parse import quote
|
||
|
||
from litellm.proxy.google_endpoints.agents_endpoints import (
|
||
list_gemini_agent_versions,
|
||
)
|
||
|
||
template = json.dumps({"api_key": "AIzaVersionsTest"})
|
||
|
||
with patch(
|
||
"litellm.proxy.google_endpoints.agents_endpoints.ProxyBaseLLMRequestProcessing"
|
||
) as MockProcessor:
|
||
instance = MockProcessor.return_value
|
||
instance.base_process_llm_request = AsyncMock(return_value=MagicMock())
|
||
|
||
request = _make_endpoint_request(f"litellm_params_template={quote(template)}")
|
||
await list_gemini_agent_versions(
|
||
request=request,
|
||
name="my-agent",
|
||
fastapi_response=MagicMock(),
|
||
user_api_key_dict=user_api_key_dict,
|
||
)
|
||
|
||
init_data = MockProcessor.call_args[1]["data"]
|
||
assert init_data.get("api_key") == "AIzaVersionsTest"
|
||
assert init_data.get("name") == "my-agent"
|
||
assert init_data.get("custom_llm_provider") == "gemini"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_gemini_agent_name_not_overwritten_by_query_param(
|
||
mock_srv, user_api_key_dict
|
||
):
|
||
"""Path-param ``name`` must not be replaced by an attacker-controlled query param."""
|
||
from urllib.parse import quote
|
||
|
||
from litellm.proxy.google_endpoints.agents_endpoints import get_gemini_agent
|
||
|
||
with patch(
|
||
"litellm.proxy.google_endpoints.agents_endpoints.ProxyBaseLLMRequestProcessing"
|
||
) as MockProcessor:
|
||
instance = MockProcessor.return_value
|
||
instance.base_process_llm_request = AsyncMock(return_value=MagicMock())
|
||
|
||
# Even if a caller tries to inject "name" via flat query param, it is
|
||
# ignored (flat params are not merged). The path-param name wins.
|
||
# ``api_key`` is supplied via the JSON template (required for non-admin
|
||
# callers — see test_*_non_admin_without_api_key_is_rejected below).
|
||
template = json.dumps({"api_key": "AIzaTest"})
|
||
request = _make_endpoint_request(
|
||
f"name=INJECTED&litellm_params_template={quote(template)}"
|
||
)
|
||
await get_gemini_agent(
|
||
request=request,
|
||
name="real-agent",
|
||
fastapi_response=MagicMock(),
|
||
user_api_key_dict=user_api_key_dict,
|
||
)
|
||
|
||
init_data = MockProcessor.call_args[1]["data"]
|
||
assert init_data["name"] == "real-agent"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_list_agents_template_via_query_param(mock_srv, user_api_key_dict):
|
||
"""litellm_params_template in query string is expanded."""
|
||
from litellm.proxy.google_endpoints.agents_endpoints import list_gemini_agents
|
||
from urllib.parse import quote
|
||
|
||
template = json.dumps({"api_key": "TemplateKey", "vertex_project": "proj-x"})
|
||
|
||
with patch(
|
||
"litellm.proxy.google_endpoints.agents_endpoints.ProxyBaseLLMRequestProcessing"
|
||
) as MockProcessor:
|
||
instance = MockProcessor.return_value
|
||
instance.base_process_llm_request = AsyncMock(return_value=MagicMock())
|
||
|
||
request = _make_endpoint_request(f"litellm_params_template={quote(template)}")
|
||
await list_gemini_agents(
|
||
request=request,
|
||
fastapi_response=MagicMock(),
|
||
user_api_key_dict=user_api_key_dict,
|
||
)
|
||
|
||
init_data = MockProcessor.call_args[1]["data"]
|
||
assert init_data["api_key"] == "TemplateKey"
|
||
assert init_data["vertex_project"] == "proj-x"
|
||
assert "litellm_params_template" not in init_data
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Security guards (veria-flagged findings)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.fixture
|
||
def proxy_admin_user_api_key_dict():
|
||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||
|
||
return UserAPIKeyAuth(
|
||
api_key="sk-admin",
|
||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_list_agents_non_admin_without_api_key_is_rejected(
|
||
mock_srv, user_api_key_dict
|
||
):
|
||
"""Non-admin callers must supply an explicit api_key — the proxy must not
|
||
silently fall back to the operator's shared GOOGLE_API_KEY/GEMINI_API_KEY.
|
||
"""
|
||
from fastapi import HTTPException
|
||
|
||
from litellm.proxy.google_endpoints.agents_endpoints import list_gemini_agents
|
||
|
||
with patch(
|
||
"litellm.proxy.google_endpoints.agents_endpoints.ProxyBaseLLMRequestProcessing"
|
||
) as MockProcessor:
|
||
instance = MockProcessor.return_value
|
||
instance.base_process_llm_request = AsyncMock(return_value=MagicMock())
|
||
|
||
request = _make_endpoint_request("")
|
||
with pytest.raises(HTTPException) as excinfo:
|
||
await list_gemini_agents(
|
||
request=request,
|
||
fastapi_response=MagicMock(),
|
||
user_api_key_dict=user_api_key_dict,
|
||
)
|
||
assert excinfo.value.status_code == 401
|
||
# Processor must never be invoked
|
||
instance.base_process_llm_request.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_delete_agent_non_admin_without_api_key_is_rejected(
|
||
mock_srv, user_api_key_dict
|
||
):
|
||
from fastapi import HTTPException
|
||
|
||
from litellm.proxy.google_endpoints.agents_endpoints import delete_gemini_agent
|
||
|
||
with patch(
|
||
"litellm.proxy.google_endpoints.agents_endpoints.ProxyBaseLLMRequestProcessing"
|
||
) as MockProcessor:
|
||
instance = MockProcessor.return_value
|
||
instance.base_process_llm_request = AsyncMock(return_value=MagicMock())
|
||
|
||
request = _make_endpoint_request("")
|
||
with pytest.raises(HTTPException) as excinfo:
|
||
await delete_gemini_agent(
|
||
request=request,
|
||
name="my-agent",
|
||
fastapi_response=MagicMock(),
|
||
user_api_key_dict=user_api_key_dict,
|
||
)
|
||
assert excinfo.value.status_code == 401
|
||
instance.base_process_llm_request.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_create_agent_non_admin_without_api_key_is_rejected(
|
||
mock_srv, user_api_key_dict
|
||
):
|
||
from fastapi import HTTPException
|
||
|
||
from litellm.proxy.google_endpoints.agents_endpoints import create_gemini_agent
|
||
|
||
with (
|
||
patch(
|
||
"litellm.proxy.google_endpoints.agents_endpoints.ProxyBaseLLMRequestProcessing"
|
||
) as MockProcessor,
|
||
patch(
|
||
"litellm.proxy.google_endpoints.agents_endpoints._read_request_body",
|
||
new=AsyncMock(return_value={"name": "agent-1", "base_agent": "waverunner"}),
|
||
),
|
||
):
|
||
instance = MockProcessor.return_value
|
||
instance.base_process_llm_request = AsyncMock(return_value=MagicMock())
|
||
|
||
request = _make_endpoint_request("")
|
||
with pytest.raises(HTTPException) as excinfo:
|
||
await create_gemini_agent(
|
||
request=request,
|
||
fastapi_response=MagicMock(),
|
||
user_api_key_dict=user_api_key_dict,
|
||
)
|
||
assert excinfo.value.status_code == 401
|
||
instance.base_process_llm_request.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_list_agents_proxy_admin_may_use_env_fallback(
|
||
mock_srv, proxy_admin_user_api_key_dict
|
||
):
|
||
"""Proxy admins (master key) keep the env-fallback convenience."""
|
||
from litellm.proxy.google_endpoints.agents_endpoints import list_gemini_agents
|
||
|
||
with patch(
|
||
"litellm.proxy.google_endpoints.agents_endpoints.ProxyBaseLLMRequestProcessing"
|
||
) as MockProcessor:
|
||
instance = MockProcessor.return_value
|
||
instance.base_process_llm_request = AsyncMock(return_value=MagicMock())
|
||
|
||
request = _make_endpoint_request("")
|
||
await list_gemini_agents(
|
||
request=request,
|
||
fastapi_response=MagicMock(),
|
||
user_api_key_dict=proxy_admin_user_api_key_dict,
|
||
)
|
||
|
||
init_data = MockProcessor.call_args[1]["data"]
|
||
assert "api_key" not in init_data
|
||
instance.base_process_llm_request.assert_awaited_once()
|
||
|
||
|
||
def test_validate_environment_rejects_api_base_override_without_explicit_key(
|
||
monkeypatch,
|
||
):
|
||
"""SECURITY: caller-supplied api_base must be paired with an explicit
|
||
api_key — otherwise the proxy's shared GOOGLE_API_KEY leaks to the
|
||
attacker-controlled host via the x-goog-api-key header.
|
||
"""
|
||
from litellm.llms.gemini.agents.transformation import GeminiAgentsConfig
|
||
|
||
# Even if env-fallback is available, api_base override must require api_key.
|
||
monkeypatch.setenv("GOOGLE_API_KEY", "AIzaSharedSecret")
|
||
|
||
cfg = GeminiAgentsConfig()
|
||
with pytest.raises(ValueError, match="api_base"):
|
||
cfg.validate_environment(
|
||
headers={},
|
||
litellm_params={"api_base": "https://attacker.example"},
|
||
)
|
||
|
||
|
||
def test_validate_environment_allows_api_base_with_explicit_key(monkeypatch):
|
||
"""api_base override is OK when paired with an explicit api_key."""
|
||
from litellm.llms.gemini.agents.transformation import GeminiAgentsConfig
|
||
|
||
monkeypatch.delenv("GOOGLE_API_KEY", raising=False)
|
||
monkeypatch.delenv("GEMINI_API_KEY", raising=False)
|
||
|
||
cfg = GeminiAgentsConfig()
|
||
headers = cfg.validate_environment(
|
||
headers={},
|
||
litellm_params={
|
||
"api_base": "https://my-gemini-proxy.example",
|
||
"api_key": "AIzaCallerOwned",
|
||
},
|
||
)
|
||
assert headers["x-goog-api-key"] == "AIzaCallerOwned"
|
||
|
||
|
||
def test_validate_environment_env_fallback_when_no_api_base_override(monkeypatch):
|
||
"""Without api_base override, env fallback continues to work for SDK use."""
|
||
from litellm.llms.gemini.agents.transformation import GeminiAgentsConfig
|
||
|
||
monkeypatch.setenv("GOOGLE_API_KEY", "AIzaFromEnv")
|
||
monkeypatch.delenv("GEMINI_API_KEY", raising=False)
|
||
|
||
cfg = GeminiAgentsConfig()
|
||
headers = cfg.validate_environment(headers={}, litellm_params={})
|
||
assert headers["x-goog-api-key"] == "AIzaFromEnv"
|