mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-13 13:10:18 +00:00
Alexsander Hamir
2e7b554747
* fix: enforce team member budget check in common_checks - Add missing team member budget validation in common_checks() function - Checks team membership budget when team key is used - Raises BudgetExceededError when team member spend exceeds max_budget_in_team - Follows same pattern as other budget checks (team, user, end_user) - Uses cached get_team_membership() for performance - Fix AttributeError in lowest_tpm_rpm.py - Add null check for model_info before accessing .get() method - Prevents 'NoneType' object has no attribute 'get' error - Add unit tests for team member budget enforcement - Test budget exceeded scenario - Test within budget scenario - Test edge cases (no budget, no membership, personal keys) - Tests run without requiring proxy server Fixes failing test: test_users_in_team_budget * fix: mock get_async_httpx_client in test_langsmith_key_based_logging - Mock get_async_httpx_client to return a mock AsyncHTTPHandler instance - Fixes test failure where mock_post was never called - LangsmithLogger creates its own httpx client instance via get_async_httpx_client, so we need to mock the factory function rather than the class method - Use MagicMock for response.raise_for_status (sync method) instead of AsyncMock * fix: resolve linting errors (PLR0915, F401) - Remove unused imports (datetime, ServiceLoggerPayload) from arize_phoenix.py - Extract health ping setup logic from RedisCache.__init__ to reduce statement count - Extract team member budget check from common_checks to reduce statement count * fix: resolve type errors in ChatCompletionToolCallChunk construction - Cast type field to Literal['function'] to satisfy TypedDict requirements - Ensure arguments field is explicitly str type to match TypedDict signature - Fixes pyright errors for incompatible types in transformation.py
365 lines
11 KiB
Python
365 lines
11 KiB
Python
"""
|
|
Unit tests for team member budget checks in common_checks.
|
|
These tests verify the team member budget enforcement without requiring a proxy server.
|
|
"""
|
|
import pytest
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
from fastapi import Request
|
|
|
|
import litellm
|
|
from litellm.proxy._types import (
|
|
LiteLLM_BudgetTable,
|
|
LiteLLM_TeamMembership,
|
|
LiteLLM_TeamTable,
|
|
LiteLLM_UserTable,
|
|
UserAPIKeyAuth,
|
|
)
|
|
from litellm.proxy.auth.auth_checks import common_checks, get_team_membership
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_team_member_budget_check_exceeds_budget():
|
|
"""Test that common_checks raises BudgetExceededError when team member spend exceeds budget."""
|
|
request_body = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
}
|
|
|
|
# Create team object
|
|
team_object = LiteLLM_TeamTable(
|
|
team_id="test-team-1",
|
|
team_alias="Test Team",
|
|
spend=0.0,
|
|
max_budget=None,
|
|
)
|
|
|
|
# Create user object
|
|
user_object = LiteLLM_UserTable(
|
|
user_id="test-user-1",
|
|
spend=0.0,
|
|
max_budget=None,
|
|
)
|
|
|
|
# Create valid token
|
|
valid_token = UserAPIKeyAuth(
|
|
token="test-token",
|
|
user_id="test-user-1",
|
|
team_id="test-team-1",
|
|
models=["gpt-3.5-turbo"],
|
|
)
|
|
|
|
# Create team membership with budget exceeded
|
|
team_membership = LiteLLM_TeamMembership(
|
|
user_id="test-user-1",
|
|
team_id="test-team-1",
|
|
spend=0.0000002, # Exceeds budget
|
|
litellm_budget_table=LiteLLM_BudgetTable(
|
|
max_budget=0.0000001, # Very small budget
|
|
),
|
|
)
|
|
|
|
mock_request = MagicMock(spec=Request)
|
|
mock_prisma_client = MagicMock()
|
|
mock_user_api_key_cache = MagicMock()
|
|
mock_proxy_logging_obj = MagicMock()
|
|
|
|
# Mock get_team_membership to return our team membership
|
|
with patch(
|
|
"litellm.proxy.auth.auth_checks.get_team_membership",
|
|
new_callable=AsyncMock,
|
|
return_value=team_membership,
|
|
), patch(
|
|
"litellm.proxy.proxy_server.prisma_client", mock_prisma_client
|
|
), patch(
|
|
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
|
|
):
|
|
# Should raise BudgetExceededError
|
|
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
|
await common_checks(
|
|
request_body=request_body,
|
|
team_object=team_object,
|
|
user_object=user_object,
|
|
end_user_object=None,
|
|
global_proxy_spend=None,
|
|
general_settings={},
|
|
route="/chat/completions",
|
|
llm_router=None,
|
|
proxy_logging_obj=mock_proxy_logging_obj,
|
|
valid_token=valid_token,
|
|
request=mock_request,
|
|
)
|
|
|
|
# Verify error message contains expected text
|
|
assert "Budget has been exceeded" in str(exc_info.value)
|
|
assert "test-user-1" in str(exc_info.value)
|
|
assert "test-team-1" in str(exc_info.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_team_member_budget_check_within_budget():
|
|
"""Test that common_checks passes when team member spend is within budget."""
|
|
request_body = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
}
|
|
|
|
# Create team object
|
|
team_object = LiteLLM_TeamTable(
|
|
team_id="test-team-1",
|
|
team_alias="Test Team",
|
|
spend=0.0,
|
|
max_budget=None,
|
|
)
|
|
|
|
# Create user object
|
|
user_object = LiteLLM_UserTable(
|
|
user_id="test-user-1",
|
|
spend=0.0,
|
|
max_budget=None,
|
|
)
|
|
|
|
# Create valid token
|
|
valid_token = UserAPIKeyAuth(
|
|
token="test-token",
|
|
user_id="test-user-1",
|
|
team_id="test-team-1",
|
|
models=["gpt-3.5-turbo"],
|
|
)
|
|
|
|
# Create team membership within budget
|
|
team_membership = LiteLLM_TeamMembership(
|
|
user_id="test-user-1",
|
|
team_id="test-team-1",
|
|
spend=0.00000005, # Within budget
|
|
litellm_budget_table=LiteLLM_BudgetTable(
|
|
max_budget=0.0000001,
|
|
),
|
|
)
|
|
|
|
mock_request = MagicMock(spec=Request)
|
|
mock_prisma_client = MagicMock()
|
|
mock_user_api_key_cache = MagicMock()
|
|
mock_proxy_logging_obj = MagicMock()
|
|
|
|
# Mock get_team_membership to return our team membership
|
|
with patch(
|
|
"litellm.proxy.auth.auth_checks.get_team_membership",
|
|
new_callable=AsyncMock,
|
|
return_value=team_membership,
|
|
), patch(
|
|
"litellm.proxy.proxy_server.prisma_client", mock_prisma_client
|
|
), patch(
|
|
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
|
|
):
|
|
# Should not raise an exception
|
|
result = await common_checks(
|
|
request_body=request_body,
|
|
team_object=team_object,
|
|
user_object=user_object,
|
|
end_user_object=None,
|
|
global_proxy_spend=None,
|
|
general_settings={},
|
|
route="/chat/completions",
|
|
llm_router=None,
|
|
proxy_logging_obj=mock_proxy_logging_obj,
|
|
valid_token=valid_token,
|
|
request=mock_request,
|
|
)
|
|
|
|
assert result is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_team_member_budget_check_no_budget_set():
|
|
"""Test that common_checks passes when team member has no budget set."""
|
|
request_body = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
}
|
|
|
|
# Create team object
|
|
team_object = LiteLLM_TeamTable(
|
|
team_id="test-team-1",
|
|
team_alias="Test Team",
|
|
spend=0.0,
|
|
max_budget=None,
|
|
)
|
|
|
|
# Create user object
|
|
user_object = LiteLLM_UserTable(
|
|
user_id="test-user-1",
|
|
spend=0.0,
|
|
max_budget=None,
|
|
)
|
|
|
|
# Create valid token
|
|
valid_token = UserAPIKeyAuth(
|
|
token="test-token",
|
|
user_id="test-user-1",
|
|
team_id="test-team-1",
|
|
models=["gpt-3.5-turbo"],
|
|
)
|
|
|
|
# Create team membership without budget
|
|
team_membership = LiteLLM_TeamMembership(
|
|
user_id="test-user-1",
|
|
team_id="test-team-1",
|
|
spend=0.0,
|
|
litellm_budget_table=None, # No budget set
|
|
)
|
|
|
|
mock_request = MagicMock(spec=Request)
|
|
mock_prisma_client = MagicMock()
|
|
mock_user_api_key_cache = MagicMock()
|
|
mock_proxy_logging_obj = MagicMock()
|
|
|
|
# Mock get_team_membership to return our team membership
|
|
with patch(
|
|
"litellm.proxy.auth.auth_checks.get_team_membership",
|
|
new_callable=AsyncMock,
|
|
return_value=team_membership,
|
|
), patch(
|
|
"litellm.proxy.proxy_server.prisma_client", mock_prisma_client
|
|
), patch(
|
|
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
|
|
):
|
|
# Should not raise an exception (no budget means no limit)
|
|
result = await common_checks(
|
|
request_body=request_body,
|
|
team_object=team_object,
|
|
user_object=user_object,
|
|
end_user_object=None,
|
|
global_proxy_spend=None,
|
|
general_settings={},
|
|
route="/chat/completions",
|
|
llm_router=None,
|
|
proxy_logging_obj=mock_proxy_logging_obj,
|
|
valid_token=valid_token,
|
|
request=mock_request,
|
|
)
|
|
|
|
assert result is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_team_member_budget_check_no_team_membership():
|
|
"""Test that common_checks passes when team membership doesn't exist."""
|
|
request_body = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
}
|
|
|
|
# Create team object
|
|
team_object = LiteLLM_TeamTable(
|
|
team_id="test-team-1",
|
|
team_alias="Test Team",
|
|
spend=0.0,
|
|
max_budget=None,
|
|
)
|
|
|
|
# Create user object
|
|
user_object = LiteLLM_UserTable(
|
|
user_id="test-user-1",
|
|
spend=0.0,
|
|
max_budget=None,
|
|
)
|
|
|
|
# Create valid token
|
|
valid_token = UserAPIKeyAuth(
|
|
token="test-token",
|
|
user_id="test-user-1",
|
|
team_id="test-team-1",
|
|
models=["gpt-3.5-turbo"],
|
|
)
|
|
|
|
mock_request = MagicMock(spec=Request)
|
|
mock_prisma_client = MagicMock()
|
|
mock_user_api_key_cache = MagicMock()
|
|
mock_proxy_logging_obj = MagicMock()
|
|
|
|
# Mock get_team_membership to return None (no membership)
|
|
with patch(
|
|
"litellm.proxy.auth.auth_checks.get_team_membership",
|
|
new_callable=AsyncMock,
|
|
return_value=None,
|
|
), patch(
|
|
"litellm.proxy.proxy_server.prisma_client", mock_prisma_client
|
|
), patch(
|
|
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
|
|
):
|
|
# Should not raise an exception (no membership means no budget check)
|
|
result = await common_checks(
|
|
request_body=request_body,
|
|
team_object=team_object,
|
|
user_object=user_object,
|
|
end_user_object=None,
|
|
global_proxy_spend=None,
|
|
general_settings={},
|
|
route="/chat/completions",
|
|
llm_router=None,
|
|
proxy_logging_obj=mock_proxy_logging_obj,
|
|
valid_token=valid_token,
|
|
request=mock_request,
|
|
)
|
|
|
|
assert result is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_team_member_budget_check_personal_key_not_team():
|
|
"""Test that team member budget check is skipped for personal keys (no team)."""
|
|
request_body = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
}
|
|
|
|
# No team object (personal key)
|
|
team_object = None
|
|
|
|
# Create user object
|
|
user_object = LiteLLM_UserTable(
|
|
user_id="test-user-1",
|
|
spend=0.0,
|
|
max_budget=None,
|
|
)
|
|
|
|
# Create valid token without team
|
|
valid_token = UserAPIKeyAuth(
|
|
token="test-token",
|
|
user_id="test-user-1",
|
|
team_id=None, # Personal key
|
|
models=["gpt-3.5-turbo"],
|
|
)
|
|
|
|
mock_request = MagicMock(spec=Request)
|
|
mock_prisma_client = MagicMock()
|
|
mock_user_api_key_cache = MagicMock()
|
|
mock_proxy_logging_obj = MagicMock()
|
|
|
|
# get_team_membership should not be called for personal keys
|
|
with patch(
|
|
"litellm.proxy.auth.auth_checks.get_team_membership",
|
|
new_callable=AsyncMock,
|
|
) as mock_get_team_membership, patch(
|
|
"litellm.proxy.proxy_server.prisma_client", mock_prisma_client
|
|
), patch(
|
|
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
|
|
):
|
|
result = await common_checks(
|
|
request_body=request_body,
|
|
team_object=team_object,
|
|
user_object=user_object,
|
|
end_user_object=None,
|
|
global_proxy_spend=None,
|
|
general_settings={},
|
|
route="/chat/completions",
|
|
llm_router=None,
|
|
proxy_logging_obj=mock_proxy_logging_obj,
|
|
valid_token=valid_token,
|
|
request=mock_request,
|
|
)
|
|
|
|
# Should pass and get_team_membership should not be called
|
|
assert result is True
|
|
mock_get_team_membership.assert_not_called()
|