mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
390 lines
16 KiB
Python
390 lines
16 KiB
Python
"""Regression tests for ProxyExtrasDBManager v2 migration resolver.
|
|
|
|
The v2 resolver is opt-in via `--use_v2_migration_resolver` / the
|
|
`use_v2_resolver=True` kwarg. These tests exercise the v2 path; the v1
|
|
(default) behavior is unchanged from pre-fix.
|
|
"""
|
|
|
|
import subprocess
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, Mock, patch
|
|
|
|
import pytest
|
|
|
|
from litellm_proxy_extras.utils import (
|
|
ProxyExtrasDBManager,
|
|
_max_migration_timestamp,
|
|
_migration_timestamp,
|
|
)
|
|
|
|
|
|
def _fake_migrate_deploy_failure(returncode: int, stderr: str):
|
|
def _run(*args, **kwargs):
|
|
raise subprocess.CalledProcessError(
|
|
returncode=returncode,
|
|
cmd=args[0],
|
|
stderr=stderr,
|
|
output="",
|
|
)
|
|
|
|
return _run
|
|
|
|
|
|
def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path):
|
|
"""v2: a permission failure during migrate deploy raises RuntimeError."""
|
|
_stub_v2_env(monkeypatch, tmp_path)
|
|
|
|
stderr = (
|
|
"Error: P3018\nMigration name: 20250326162113_baseline\n"
|
|
"Database error code: 42501\npermission denied for schema public"
|
|
)
|
|
with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
|
with pytest.raises(RuntimeError, match="permission"):
|
|
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
|
|
|
|
|
def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path):
|
|
"""v2: a non-idempotent migration failure raises (no silent recovery)."""
|
|
_stub_v2_env(monkeypatch, tmp_path)
|
|
|
|
stderr = (
|
|
"Error: P3009\nMigration `20260101000000_genuinely_broken` failed\n"
|
|
'Reason: syntax error at or near "BRKN" LINE 42'
|
|
)
|
|
with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
|
with pytest.raises(RuntimeError, match="Migration completion could not be verified"):
|
|
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
|
|
|
|
|
def test_strip_prisma_query_params_removes_connection_limit():
|
|
"""DATABASE_URLs with Prisma-specific params should be parseable by psycopg."""
|
|
url = "postgresql://u:p@h:5432/db?connection_limit=100&pool_timeout=60&sslmode=require"
|
|
stripped = ProxyExtrasDBManager._strip_prisma_query_params(url)
|
|
assert "connection_limit" not in stripped
|
|
assert "pool_timeout" not in stripped
|
|
assert "sslmode=require" in stripped
|
|
|
|
|
|
def test_strip_prisma_query_params_passthrough_no_query():
|
|
"""URLs without query strings are returned unchanged."""
|
|
url = "postgresql://u:p@h:5432/db"
|
|
assert ProxyExtrasDBManager._strip_prisma_query_params(url) == url
|
|
|
|
|
|
def test_migration_timestamp_extracts_leading_digits():
|
|
assert _migration_timestamp("20260101000000_add_foo") == 20260101000000
|
|
assert _migration_timestamp("20250326162113_baseline") == 20250326162113
|
|
|
|
|
|
def test_migration_timestamp_returns_zero_on_malformed():
|
|
assert _migration_timestamp("0_init") == 0
|
|
assert _migration_timestamp("not_a_migration") == 0
|
|
|
|
|
|
def test_max_migration_timestamp():
|
|
names = {"20250326000000_a", "20260415000000_b", "20251115000000_c"}
|
|
assert _max_migration_timestamp(names) == 20260415000000
|
|
|
|
|
|
def test_max_migration_timestamp_empty_set():
|
|
assert _max_migration_timestamp(set()) == 0
|
|
|
|
|
|
def test_v1_default_still_calls_resolve_all_migrations(monkeypatch, tmp_path):
|
|
"""v1 (default) continues to call _resolve_all_migrations on the happy path.
|
|
|
|
This is the existing buggy behavior — we're not fixing it in v1, only
|
|
offering v2 as opt-in. This test pins the default so that a future
|
|
inadvertent default flip is caught.
|
|
"""
|
|
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
|
(tmp_path / "schema.prisma").write_text("// stub")
|
|
|
|
# Stub `prisma migrate deploy` to claim success with pending migrations
|
|
# applied, which is the code path that triggers the legacy post-migration
|
|
# sanity check (a call to _resolve_all_migrations).
|
|
class FakeResult:
|
|
stdout = "Applied migration.\n"
|
|
stderr = ""
|
|
|
|
def fake_run(cmd, *args, **kwargs):
|
|
return FakeResult()
|
|
|
|
resolve_called = {"n": 0}
|
|
|
|
def fake_resolve(*args, **kwargs):
|
|
resolve_called["n"] += 1
|
|
|
|
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", fake_run)
|
|
monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", fake_resolve)
|
|
|
|
ok = ProxyExtrasDBManager.setup_database(use_migrate=True) # v2 flag NOT set
|
|
assert ok is True
|
|
assert resolve_called["n"] == 1, "v1 default should still invoke the legacy path"
|
|
|
|
|
|
def test_v2_db_push_wraps_subprocess_error_as_runtime_error(monkeypatch, tmp_path):
|
|
"""v2: a failing `prisma db push` must raise RuntimeError, not leak
|
|
CalledProcessError past proxy_cli.py's `except RuntimeError`."""
|
|
monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(tmp_path))
|
|
|
|
stderr = "db push error"
|
|
with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
|
with pytest.raises(RuntimeError, match="prisma db push failed"):
|
|
ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True)
|
|
|
|
|
|
def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path):
|
|
"""_warn_if_db_ahead_of_head must never raise — it's informational.
|
|
|
|
Non-connection DB errors (e.g. InsufficientPrivilege from a user
|
|
without SELECT on _prisma_migrations) must be caught, not propagated.
|
|
"""
|
|
import psycopg
|
|
|
|
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
|
monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(tmp_path))
|
|
|
|
class _FakeConn:
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *a):
|
|
return False
|
|
|
|
def execute(self, *a, **kw):
|
|
# Simulate an InsufficientPrivilege (subclass of DatabaseError).
|
|
raise psycopg.errors.InsufficientPrivilege("permission denied")
|
|
|
|
def _fake_connect(*a, **kw):
|
|
return _FakeConn()
|
|
|
|
monkeypatch.setattr("psycopg.connect", _fake_connect)
|
|
|
|
# Must not raise.
|
|
ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path))
|
|
|
|
|
|
def test_v2_duplicate_object_p3009_is_not_marked_applied(monkeypatch, tmp_path):
|
|
_stub_v2_env(monkeypatch, tmp_path, ledger_logs="relation already exists")
|
|
stderr = "Error: P3009\nMigration `20260101000000_some_migration` failed\nrelation already exists"
|
|
with patch(
|
|
"litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)
|
|
) as run:
|
|
with pytest.raises(RuntimeError, match="Migration completion could not be verified"):
|
|
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
|
assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == (
|
|
["migrate", "deploy"],
|
|
)
|
|
|
|
|
|
def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path):
|
|
_stub_v2_env(monkeypatch, tmp_path)
|
|
run = Mock(side_effect=_succeed_after(0, ""))
|
|
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", run)
|
|
|
|
assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) is True
|
|
assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == (
|
|
["migrate", "deploy"],
|
|
)
|
|
|
|
|
|
_DEADLOCK_P3018_STDERR = (
|
|
"Error: P3018\n"
|
|
"Migration name: 20260415120000_health_check_latest_per_model_index\n"
|
|
"Database error code: 40P01\n"
|
|
"deadlock detected"
|
|
)
|
|
|
|
|
|
def _stub_v2_env(monkeypatch, tmp_path, ledger_logs=""):
|
|
import psycopg
|
|
|
|
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
|
monkeypatch.delenv("DIRECT_URL", raising=False)
|
|
monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(tmp_path))
|
|
monkeypatch.setattr("time.sleep", lambda _: None)
|
|
connection = MagicMock()
|
|
connection.__enter__.return_value = connection
|
|
cursor = connection.cursor.return_value.__enter__.return_value
|
|
cursor.execute.return_value = cursor
|
|
cursor.fetchone.return_value = SimpleNamespace(acquired=True)
|
|
cursor.fetchall.return_value = []
|
|
empty = MagicMock()
|
|
empty.fetchall.return_value = []
|
|
empty.fetchone.return_value = None
|
|
ledger = MagicMock()
|
|
ledger.fetchone.return_value = (ledger_logs,)
|
|
|
|
def execute(query, *args, **kwargs):
|
|
if "SELECT logs FROM" in str(query):
|
|
if ledger_logs is None:
|
|
raise psycopg.OperationalError("ledger is unavailable")
|
|
return ledger
|
|
return empty
|
|
|
|
connection.execute.side_effect = execute
|
|
monkeypatch.setattr("psycopg.connect", lambda *args, **kwargs: connection)
|
|
|
|
|
|
def _succeed_after(failures: int, stderr: str):
|
|
calls = {"n": 0}
|
|
|
|
class _OkResult:
|
|
stdout = "Applied migration.\n"
|
|
stderr = ""
|
|
|
|
def _run(*args, **kwargs):
|
|
if "deploy" not in args[0]:
|
|
return _OkResult()
|
|
calls["n"] += 1
|
|
if calls["n"] <= failures:
|
|
raise subprocess.CalledProcessError(returncode=1, cmd=args[0], stderr=stderr, output="")
|
|
return _OkResult()
|
|
|
|
return _run
|
|
|
|
|
|
def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path):
|
|
"""v2: losing the migrate deploy deadlock race against a concurrent
|
|
instance rolls the ledger row back and retries instead of dying."""
|
|
_stub_v2_env(monkeypatch, tmp_path)
|
|
|
|
run = Mock(side_effect=_succeed_after(1, _DEADLOCK_P3018_STDERR))
|
|
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", run)
|
|
|
|
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
|
assert ok is True
|
|
assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == (
|
|
["migrate", "deploy"],
|
|
["migrate", "resolve", "--rolled-back", "20260415120000_health_check_latest_per_model_index"],
|
|
["migrate", "deploy"],
|
|
)
|
|
|
|
|
|
def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path):
|
|
"""v2: a deadlock on every attempt still fails after the retry budget."""
|
|
_stub_v2_env(monkeypatch, tmp_path)
|
|
|
|
with patch(
|
|
"litellm_proxy_extras.prisma_toolchain.run_prisma",
|
|
side_effect=_fake_migrate_deploy_failure(1, _DEADLOCK_P3018_STDERR),
|
|
):
|
|
with pytest.raises(RuntimeError, match="after 4 attempts"):
|
|
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
|
|
|
|
|
def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_path):
|
|
"""v2: the surviving instance sees the victim's failed ledger row as P3009.
|
|
When that row's logs show a deadlock, roll it back and retry."""
|
|
_stub_v2_env(monkeypatch, tmp_path, ledger_logs="ERROR: deadlock detected\nDETAIL: Process 72 waits for ShareLock")
|
|
|
|
stderr = (
|
|
"Error: P3009\n"
|
|
"migrate found failed migrations in the target database\n"
|
|
"The `20260415120000_health_check_latest_per_model_index` migration "
|
|
"started at 2026-09-01 18:46:13 UTC failed"
|
|
)
|
|
run = Mock(side_effect=_succeed_after(1, stderr))
|
|
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", run)
|
|
|
|
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
|
assert ok is True
|
|
assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == (
|
|
["migrate", "deploy"],
|
|
["migrate", "resolve", "--rolled-back", "20260415120000_health_check_latest_per_model_index"],
|
|
["migrate", "deploy"],
|
|
)
|
|
|
|
|
|
def test_v2_p3009_empty_ledger_logs_do_not_prove_completion(monkeypatch, tmp_path):
|
|
_stub_v2_env(monkeypatch, tmp_path)
|
|
stderr = (
|
|
"Error: P3009\n"
|
|
"migrate found failed migrations in the target database\n"
|
|
"The `20260415120000_health_check_latest_per_model_index` migration "
|
|
"started at 2026-09-01 18:46:13 UTC failed"
|
|
)
|
|
with patch(
|
|
"litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)
|
|
) as run:
|
|
with pytest.raises(RuntimeError, match="Migration completion could not be verified"):
|
|
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
|
assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == (
|
|
["migrate", "deploy"],
|
|
)
|
|
|
|
|
|
def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path):
|
|
"""v2: an unreadable ledger cannot establish that P3009 was a deadlock."""
|
|
_stub_v2_env(monkeypatch, tmp_path, ledger_logs=None)
|
|
|
|
stderr = (
|
|
"Error: P3009\n"
|
|
"migrate found failed migrations in the target database\n"
|
|
"The `20260415120000_health_check_latest_per_model_index` migration "
|
|
"started at 2026-09-01 18:46:13 UTC failed"
|
|
)
|
|
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr))
|
|
|
|
with pytest.raises(RuntimeError, match="Migration completion could not be verified"):
|
|
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
|
|
|
|
|
def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path):
|
|
"""v2: a failed ledger row whose logs show a real SQL error stays fatal."""
|
|
_stub_v2_env(monkeypatch, tmp_path, ledger_logs='ERROR: syntax error at or near "BRKN"')
|
|
|
|
stderr = (
|
|
"Error: P3009\n"
|
|
"migrate found failed migrations in the target database\n"
|
|
"The `20260101000000_genuinely_broken` migration started at "
|
|
"2026-09-01 18:46:13 UTC failed"
|
|
)
|
|
|
|
with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
|
with pytest.raises(RuntimeError, match="Migration completion could not be verified"):
|
|
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
|
|
|
|
|
def test_v2_bare_deadlock_stderr_retries(monkeypatch, tmp_path):
|
|
"""v2: a deadlock reported without a Prisma error code (the advisory-lock
|
|
waiter as victim) is retried, not fatal."""
|
|
_stub_v2_env(monkeypatch, tmp_path)
|
|
monkeypatch.setattr(
|
|
"litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, "Database error: deadlock detected")
|
|
)
|
|
|
|
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
|
assert ok is True
|
|
|
|
|
|
_P1002_ADVISORY_LOCK_STDERR = (
|
|
"Error: P1002\n\n"
|
|
"The database server at `127.0.0.1`:`45743` was reached but timed out.\n\n"
|
|
"Context: Timed out trying to acquire a postgres advisory lock "
|
|
"(SELECT pg_advisory_lock(72707369)). Elapsed: 10000ms."
|
|
)
|
|
|
|
|
|
def test_v2_advisory_lock_timeout_retries(monkeypatch, tmp_path):
|
|
"""v2: the advisory-lock waiter that times out while a peer's retry holds
|
|
the lock retries instead of dying."""
|
|
_stub_v2_env(monkeypatch, tmp_path)
|
|
monkeypatch.setattr(
|
|
"litellm_proxy_extras.prisma_toolchain.run_prisma",
|
|
_succeed_after(2, _P1002_ADVISORY_LOCK_STDERR),
|
|
)
|
|
|
|
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
|
assert ok is True
|
|
|
|
|
|
def test_v2_p1002_without_advisory_lock_context_still_raises(monkeypatch, tmp_path):
|
|
"""v2: a plain P1002 (database unreachable) stays fatal."""
|
|
_stub_v2_env(monkeypatch, tmp_path)
|
|
stderr = "Error: P1002\n\nThe database server at `db`:`5432` was reached but timed out."
|
|
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr))
|
|
|
|
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
|
|
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|