test: keep the resolver tests inside the test-quality ceilings

The moved fail-fast test carried a sys.path.insert that the uv workspace
makes unnecessary, and one pre-existing case asserted nothing beyond "did
not raise", so it could not tell a swallowed error from a skipped query.
Give it a liveness gate on the connect count instead.

Fold the resolver default/opt-out matrix into the existing db-push flag
test rather than standing up another patched test, so the flag pair, the
env var, and their precedence are covered without new mock scaffolding.
This commit is contained in:
Yuneng Jiang 2026-08-28 01:52:15 -07:00
parent fbd1339993
commit c3fc86869d
No known key found for this signature in database
2 changed files with 38 additions and 83 deletions

View file

@ -4,20 +4,11 @@ v2 is the proxy CLI default; v1 stays reachable via the `use_v2_resolver`
kwarg, which still defaults to False for direct callers.
"""
import os
import subprocess
import sys
from unittest.mock import patch
import pytest
sys.path.insert(
0,
os.path.abspath(
os.path.join(os.path.dirname(__file__), "../../litellm-proxy-extras")
),
)
from litellm_proxy_extras.utils import (
ProxyExtrasDBManager,
_max_migration_timestamp,
@ -175,13 +166,16 @@ def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path):
# Simulate an InsufficientPrivilege (subclass of DatabaseError).
raise psycopg.errors.InsufficientPrivilege("permission denied")
connects = {"n": 0}
def _fake_connect(*a, **kw):
connects["n"] += 1
return _FakeConn()
monkeypatch.setattr("psycopg.connect", _fake_connect)
# Must not raise.
ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path))
assert ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path)) is None
assert connects["n"] == 1, "the failing query must actually have been reached"
def test_v2_resolve_specific_migration_failure_raises_runtime_error(

View file

@ -1805,6 +1805,39 @@ class TestRunServerDbSetup:
use_migrate=False, use_v2_resolver=True
)
# Test 3+: the resolver default and both routes back to v1. The flag
# covers a CLI boot; USE_V2_MIGRATION_RESOLVER covers deploys that
# cannot pass one, where prisma_migration.py fixes the argv. An
# explicit flag beats the env var.
for argv, env_value, expected_v2 in (
([], None, True),
(["--use_v2_migration_resolver"], None, True),
(["--use_legacy_migration_resolver"], None, False),
([], "false", False),
([], "true", True),
(["--use_v2_migration_resolver"], "false", True),
(["--use_legacy_migration_resolver"], "true", False),
):
mock_setup_database.reset_mock()
mock_should_update_schema.reset_mock()
mock_should_update_schema.return_value = True
resolver_env = (
{"USE_V2_MIGRATION_RESOLVER": env_value}
if env_value is not None
else {}
)
os.environ.pop("USE_V2_MIGRATION_RESOLVER", None)
with patch.dict(os.environ, resolver_env):
run_server.main(
["--local", "--skip_server_startup", *argv],
standalone_mode=False,
)
assert mock_setup_database.call_args.kwargs == {
"use_migrate": True,
"use_v2_resolver": expected_v2,
}, f"argv={argv} env={env_value}"
@patch("subprocess.run")
@patch("atexit.register")
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
@ -1981,78 +2014,6 @@ class TestRunServerDbSetup:
use_migrate=True, use_v2_resolver=True
)
@pytest.mark.parametrize(
"argv, env_value, expected_v2",
[
([], None, True),
(["--use_v2_migration_resolver"], None, True),
(["--use_legacy_migration_resolver"], None, False),
([], "false", False),
([], "true", True),
(["--use_v2_migration_resolver"], "false", True),
],
)
@patch("subprocess.run")
@patch("atexit.register")
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")
def test_migration_resolver_default_and_opt_out(
self,
mock_should_update_schema,
mock_check_schema_diff,
mock_setup_database,
mock_atexit_register,
mock_subprocess_run,
argv,
env_value,
expected_v2,
):
"""The proxy defaults to v2, and both v1 opt-out routes work: the
flag for a CLI boot, USE_V2_MIGRATION_RESOLVER=false for deploys that
cannot pass one. An explicit flag beats the env var."""
from litellm.proxy.proxy_cli import run_server
mock_subprocess_run.return_value = MagicMock(returncode=0)
mock_should_update_schema.return_value = True
mock_setup_database.return_value = True
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
clean_env = {
k: v
for k, v in os.environ.items()
if k
not in ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER")
}
clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
if env_value is not None:
clean_env["USE_V2_MIGRATION_RESOLVER"] = env_value
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"sys.modules",
{
"proxy_server": mock_proxy_module,
"litellm.proxy.proxy_server": mock_proxy_module,
},
),
):
run_server.main(
["--local", "--skip_server_startup", *argv], standalone_mode=False
)
mock_setup_database.assert_called_once_with(
use_migrate=True, use_v2_resolver=expected_v2
)
# --- Module-level helpers for worker startup hook tests ---
_dummy_hook_called = False