mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
260ff5f491
A team can now carry a per-model budget map that every key on the team
inherits. A key's own model_max_budget entry for the same model takes
precedence, so it is gated on and billed to the key alone.
Backend: NewTeamRequest/UpdateTeamRequest accept model_max_budget (validated
like the key-level field, enterprise gated); the value is hydrated onto
UserAPIKeyAuth via the token view, TeamGrants and the carried budget state;
_check_team_model_budget enforces it in the centralized common checks; the
limiter meters spend under team_model_spend:<team>:<model>:<duration> and
skips the team counter when the key overrides; /team/update lets only a
proxy admin raise, re-window or drop a cap; /team/info exposes usage.
The Anthropic context-management compaction summary subrequest runs the
same team gate. Both fallback token-view SQL definitions project the column.
UI: team create and edit forms reuse the key-level ModelMaxBudgetEditor,
premium gated, sending {} to clear and omitting unchanged fields.
A key entry overrides the team cap only when it spend-gates the model
(non-negative max_budget); a row that only carries tpm/rpm limits or a
negative cap leaves the team cap in force.
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
255 lines
8.9 KiB
Python
255 lines
8.9 KiB
Python
"""
|
|
Tests for create_missing_views exception handling fix.
|
|
|
|
Verifies that real DB errors (auth failures, connection errors, etc.)
|
|
are re-raised instead of being silently swallowed, while genuine
|
|
"view not found" errors still trigger view creation.
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import AsyncMock, MagicMock, call
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_views_reraises_connection_error():
|
|
"""should re-raise exceptions that are NOT 'does not exist' errors (e.g. connection errors)."""
|
|
from litellm.proxy.db.create_views import create_missing_views
|
|
|
|
mock_db = MagicMock()
|
|
mock_db.query_raw = AsyncMock(
|
|
side_effect=Exception("connection refused: unable to connect to database")
|
|
)
|
|
mock_db.execute_raw = AsyncMock()
|
|
|
|
with pytest.raises(Exception, match="connection refused"):
|
|
await create_missing_views(mock_db)
|
|
|
|
mock_db.execute_raw.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_views_reraises_permission_error():
|
|
"""should re-raise permission denied errors, not treat them as missing views."""
|
|
from litellm.proxy.db.create_views import create_missing_views
|
|
|
|
mock_db = MagicMock()
|
|
mock_db.query_raw = AsyncMock(
|
|
side_effect=Exception(
|
|
"permission denied for table LiteLLM_VerificationTokenView"
|
|
)
|
|
)
|
|
mock_db.execute_raw = AsyncMock()
|
|
|
|
with pytest.raises(Exception, match="permission denied"):
|
|
await create_missing_views(mock_db)
|
|
|
|
mock_db.execute_raw.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_views_creates_view_on_does_not_exist():
|
|
"""should call execute_raw to create view when error contains 'does not exist'."""
|
|
from litellm.proxy.db.create_views import create_missing_views
|
|
|
|
mock_db = MagicMock()
|
|
mock_db.query_raw = AsyncMock(
|
|
side_effect=[
|
|
Exception('relation "LiteLLM_VerificationTokenView" does not exist'),
|
|
None, # MonthlyGlobalSpend exists
|
|
None, # Last30dKeysBySpend exists
|
|
None, # Last30dModelsBySpend exists
|
|
None, # MonthlyGlobalSpendPerKey exists
|
|
None, # MonthlyGlobalSpendPerUserPerKey exists
|
|
None, # DailyTagSpend exists
|
|
None, # Last30dTopEndUsersSpend exists
|
|
]
|
|
)
|
|
mock_db.execute_raw = AsyncMock(return_value=None)
|
|
|
|
await create_missing_views(mock_db)
|
|
|
|
mock_db.execute_raw.assert_called_once()
|
|
created_sql = mock_db.execute_raw.call_args[0][0]
|
|
assert 'CREATE VIEW "LiteLLM_VerificationTokenView"' in created_sql
|
|
assert "t.model_max_budget AS team_model_max_budget" in created_sql
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_views_creates_view_on_undefined_error():
|
|
"""should treat 'undefined' errors as 'view not found' and attempt creation."""
|
|
from litellm.proxy.db.create_views import create_missing_views
|
|
|
|
mock_db = MagicMock()
|
|
mock_db.query_raw = AsyncMock(
|
|
side_effect=[
|
|
Exception("undefined table LiteLLM_VerificationTokenView"),
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
]
|
|
)
|
|
mock_db.execute_raw = AsyncMock(return_value=None)
|
|
|
|
await create_missing_views(mock_db)
|
|
|
|
mock_db.execute_raw.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_views_skips_creation_when_view_exists():
|
|
"""should not call execute_raw when all views already exist."""
|
|
from litellm.proxy.db.create_views import create_missing_views
|
|
|
|
mock_db = MagicMock()
|
|
mock_db.query_raw = AsyncMock(return_value=[{"?column?": 1}])
|
|
mock_db.execute_raw = AsyncMock()
|
|
|
|
await create_missing_views(mock_db)
|
|
|
|
mock_db.execute_raw.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_views_reraises_undefined_function_error():
|
|
"""should re-raise 'undefined function' errors — bare 'undefined' is too broad
|
|
and would previously misclassify DB function errors as missing-view signals."""
|
|
from litellm.proxy.db.create_views import create_missing_views
|
|
|
|
mock_db = MagicMock()
|
|
mock_db.query_raw = AsyncMock(
|
|
side_effect=Exception("ERROR: undefined function pg_get_viewdef()")
|
|
)
|
|
mock_db.execute_raw = AsyncMock()
|
|
|
|
with pytest.raises(Exception, match="undefined function"):
|
|
await create_missing_views(mock_db)
|
|
|
|
mock_db.execute_raw.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_create_missing_views_reltuples_zero():
|
|
"""should return True when reltuples is 0 (fresh empty table)."""
|
|
from litellm.proxy.db.create_views import should_create_missing_views
|
|
|
|
mock_db = MagicMock()
|
|
mock_db.query_raw = AsyncMock(return_value=[{"reltuples": 0}])
|
|
|
|
result = await should_create_missing_views(mock_db)
|
|
assert result is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_create_missing_views_reltuples_negative_one():
|
|
"""should return True when reltuples is -1 (table created, no ANALYZE yet)."""
|
|
from litellm.proxy.db.create_views import should_create_missing_views
|
|
|
|
mock_db = MagicMock()
|
|
mock_db.query_raw = AsyncMock(return_value=[{"reltuples": -1}])
|
|
|
|
result = await should_create_missing_views(mock_db)
|
|
assert result is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_create_missing_views_reltuples_positive():
|
|
"""should return False when reltuples > 0 (table has data)."""
|
|
from litellm.proxy.db.create_views import should_create_missing_views
|
|
|
|
mock_db = MagicMock()
|
|
mock_db.query_raw = AsyncMock(return_value=[{"reltuples": 1000}])
|
|
|
|
result = await should_create_missing_views(mock_db)
|
|
assert result is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_views_creates_view_on_undefined_table_error():
|
|
"""should treat 'undefined table' as a missing-view signal and attempt creation."""
|
|
from litellm.proxy.db.create_views import create_missing_views
|
|
|
|
mock_db = MagicMock()
|
|
mock_db.query_raw = AsyncMock(
|
|
side_effect=[
|
|
Exception('undefined table "LiteLLM_VerificationTokenView"'),
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
]
|
|
)
|
|
mock_db.execute_raw = AsyncMock(return_value=None)
|
|
|
|
await create_missing_views(mock_db)
|
|
|
|
mock_db.execute_raw.assert_called_once()
|
|
|
|
|
|
# Every view create_missing_views is responsible for. Hard-coded rather than
|
|
# derived from the module, so adding a view without guarding it fails here.
|
|
EXPECTED_VIEW_COUNT = 8
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_views_tolerates_a_concurrent_creator_on_every_view():
|
|
"""A replica that loses the CREATE race must attempt every view regardless.
|
|
|
|
Regression: two proxy pods booting on a fresh DB both see every view as
|
|
absent and both issue the CREATE, and Postgres fails the loser with a
|
|
duplicate-object error on whichever views the winner got to first. Any
|
|
creation site still calling execute_raw unguarded re-raises that error and
|
|
aborts the rest of the function.
|
|
|
|
Every CREATE loses here, which is what pins the guard to all of them: an
|
|
earlier version of this fix converted only the first and the last site and
|
|
still died on MonthlyGlobalSpend against a real Postgres. Counting the
|
|
attempts is the assertion, because a partial fix simply stops early.
|
|
"""
|
|
from litellm.proxy.db.create_views import create_missing_views
|
|
|
|
mock_db = MagicMock()
|
|
mock_db.query_raw = AsyncMock(side_effect=Exception("relation does not exist"))
|
|
mock_db.execute_raw = AsyncMock(
|
|
side_effect=Exception('relation "some_view" already exists')
|
|
)
|
|
|
|
await create_missing_views(mock_db)
|
|
|
|
assert mock_db.execute_raw.await_count == EXPECTED_VIEW_COUNT, (
|
|
f"every view must still be attempted when the replica loses every race; "
|
|
f"got {mock_db.execute_raw.await_count} of {EXPECTED_VIEW_COUNT}, so a "
|
|
f"creation site is still unguarded and aborted the rest"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_views_reraises_genuine_ddl_error():
|
|
"""An already-exists guard must not swallow real DDL failures."""
|
|
from litellm.proxy.db.create_views import create_missing_views
|
|
|
|
mock_db = MagicMock()
|
|
mock_db.query_raw = AsyncMock(side_effect=Exception("relation does not exist"))
|
|
mock_db.execute_raw = AsyncMock(side_effect=Exception("syntax error at or near"))
|
|
|
|
with pytest.raises(Exception, match="syntax error"):
|
|
await create_missing_views(mock_db)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_view_tolerating_race_swallows_only_already_exists():
|
|
from litellm.proxy.db.create_views import create_view_tolerating_race
|
|
|
|
mock_db = MagicMock()
|
|
mock_db.execute_raw = AsyncMock(side_effect=Exception("duplicate object"))
|
|
await create_view_tolerating_race(mock_db, "SomeView", "CREATE VIEW ...")
|
|
|
|
mock_db.execute_raw = AsyncMock(side_effect=Exception("permission denied"))
|
|
with pytest.raises(Exception, match="permission denied"):
|
|
await create_view_tolerating_race(mock_db, "SomeView", "CREATE VIEW ...")
|