mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
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:
parent
fbd1339993
commit
c3fc86869d
2 changed files with 38 additions and 83 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue