mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
726292720c
Against a real Postgres the previous commit still died on MonthlyGlobalSpend: only 2 of the 8 creation sites went through the tolerant helper, so the losing replica re-raised on the first unguarded one and skipped the rest. The regression test now makes every CREATE lose the race and asserts all 8 are still attempted, which fails on the partial fix.
254 lines
8.8 KiB
Python
254 lines
8.8 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
|
|
|
|
|
|
@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 ...")
|