mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
25109a523b
* test(proxy): move utils, agent_endpoints and endpoint tests into tests/unit/proxy Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): package moved unit test directories Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): exclude proxy-db-owned files from the misc target Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): drop the redundant fixture docstrings in the proxy conftest Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yuneng <yuneng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
96 lines
3.3 KiB
Python
96 lines
3.3 KiB
Python
"""
|
|
Tests for RBAC enforcement on agent endpoints.
|
|
|
|
Verifies that check_feature_access_for_user is called and that a 403 is
|
|
raised when agents are disabled for internal users.
|
|
"""
|
|
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
|
|
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
|
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
|
|
RestrictedAgentAccess,
|
|
UnrestrictedAgentAccess,
|
|
)
|
|
|
|
|
|
def _make_internal_user(user_id: str = "user-1") -> UserAPIKeyAuth:
|
|
return UserAPIKeyAuth(
|
|
user_role=LitellmUserRoles.INTERNAL_USER.value,
|
|
user_id=user_id,
|
|
)
|
|
|
|
|
|
def _make_admin_user(user_id: str = "admin-1") -> UserAPIKeyAuth:
|
|
return UserAPIKeyAuth(
|
|
user_role=LitellmUserRoles.PROXY_ADMIN.value,
|
|
user_id=user_id,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# get_agents
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_agents_blocked_for_internal_user_when_disabled():
|
|
"""get_agents should raise 403 when agents are disabled for internal users."""
|
|
from litellm.proxy.agent_endpoints.endpoints import get_agents
|
|
|
|
user = _make_internal_user()
|
|
gs = {
|
|
"disable_agents_for_internal_users": True,
|
|
"allow_agents_for_team_admins": False,
|
|
}
|
|
|
|
request_mock = MagicMock()
|
|
with patch.dict("litellm.proxy.proxy_server.general_settings", gs, clear=True):
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await get_agents(request=request_mock, user_api_key_dict=user)
|
|
assert exc_info.value.status_code == 403
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_agents_allowed_when_not_disabled():
|
|
"""get_agents should not raise RBAC 403 when agents are not disabled."""
|
|
from litellm.proxy.agent_endpoints.endpoints import get_agents
|
|
|
|
user = _make_internal_user()
|
|
request_mock = MagicMock()
|
|
|
|
with patch.dict("litellm.proxy.proxy_server.general_settings", {}, clear=True):
|
|
with patch(
|
|
"litellm.proxy.agent_endpoints.agent_registry.global_agent_registry",
|
|
MagicMock(get_agent_list=MagicMock(return_value=[])),
|
|
):
|
|
with patch(
|
|
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access",
|
|
new=AsyncMock(return_value=UnrestrictedAgentAccess()),
|
|
):
|
|
result = await get_agents(request=request_mock, user_api_key_dict=user)
|
|
assert result == []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# get_agent_daily_activity
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_agent_daily_activity_blocked_when_disabled():
|
|
from litellm.proxy.agent_endpoints.endpoints import get_agent_daily_activity
|
|
|
|
user = _make_internal_user()
|
|
gs = {
|
|
"disable_agents_for_internal_users": True,
|
|
"allow_agents_for_team_admins": False,
|
|
}
|
|
|
|
with patch.dict("litellm.proxy.proxy_server.general_settings", gs, clear=True):
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await get_agent_daily_activity(user_api_key_dict=user)
|
|
assert exc_info.value.status_code == 403
|