mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
test: isolate migration CI selection and exercise resolver boundaries
This commit is contained in:
Yuneng Jiang
parent
c44757fc01
commit
a37f0b4f54
11 changed files with 93 additions and 124 deletions
2
.github/e2e-stack/select_tests.py
vendored
2
.github/e2e-stack/select_tests.py
vendored
|
|
@ -4,7 +4,7 @@ from typing import Final
|
|||
|
||||
SELECTABLE: Final = re.compile(r"^tests/e2e/([A-Za-z0-9_.-]+/)*test_[A-Za-z0-9_.-]+\.py$")
|
||||
UNSUPPORTED: Final = re.compile(
|
||||
r"^tests/e2e/(ui|claude_code|load)/"
|
||||
r"^tests/e2e/(ui|claude_code|load|migrations)/"
|
||||
r"|^tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e\.py$"
|
||||
r"|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$"
|
||||
r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$"
|
||||
|
|
|
|||
|
|
@ -859,7 +859,11 @@ class ProxyExtrasDBManager:
|
|||
def baseline_existing(migrations_dir: str) -> None:
|
||||
with migration_lock(lock_url) as coordinator:
|
||||
baseline_current_schema(
|
||||
coordinator, schema, Path(migrations_dir), _get_prisma_command(), migration_environment(_get_prisma_env())
|
||||
coordinator,
|
||||
schema,
|
||||
Path(migrations_dir),
|
||||
_get_prisma_command(),
|
||||
migration_environment(_get_prisma_env()),
|
||||
)
|
||||
|
||||
while not ProxyExtrasDBManager._run_database_v2(True, recover_completed, baseline_existing):
|
||||
|
|
@ -1054,7 +1058,8 @@ class ProxyExtrasDBManager:
|
|||
migration_name = ProxyExtrasDBManager._v2_failed_migration_name(stderr)
|
||||
if migration_name and _MIGRATION_DEADLOCK_MARKER in stderr:
|
||||
logger.info(
|
||||
"Migration %s deadlocked against a concurrent migrate deploy, rolling its ledger row back and retrying",
|
||||
"Migration %s deadlocked against a concurrent migrate deploy, "
|
||||
"rolling its ledger row back and retrying",
|
||||
migration_name,
|
||||
)
|
||||
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
|
||||
|
|
|
|||
|
|
@ -6,7 +6,8 @@ The v2 resolver is opt-in via `--use_v2_migration_resolver` / the
|
|||
"""
|
||||
|
||||
import subprocess
|
||||
from unittest.mock import patch
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -31,10 +32,7 @@ def _fake_migrate_deploy_failure(returncode: int, stderr: str):
|
|||
|
||||
def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path):
|
||||
"""v2: a permission failure during migrate deploy raises RuntimeError."""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
_stub_v2_env(monkeypatch, tmp_path)
|
||||
|
||||
stderr = (
|
||||
"Error: P3018\nMigration name: 20250326162113_baseline\n"
|
||||
|
|
@ -47,10 +45,7 @@ def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path):
|
|||
|
||||
def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path):
|
||||
"""v2: a non-idempotent migration failure raises (no silent recovery)."""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
_stub_v2_env(monkeypatch, tmp_path)
|
||||
|
||||
stderr = (
|
||||
"Error: P3009\nMigration `20260101000000_genuinely_broken` failed\n"
|
||||
|
|
@ -131,8 +126,7 @@ def test_v1_default_still_calls_resolve_all_migrations(monkeypatch, tmp_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.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
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)):
|
||||
|
|
@ -149,8 +143,7 @@ def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path):
|
|||
import psycopg
|
||||
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(tmp_path))
|
||||
|
||||
class _FakeConn:
|
||||
def __enter__(self):
|
||||
|
|
@ -173,18 +166,7 @@ def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path):
|
|||
|
||||
|
||||
def test_v2_duplicate_object_p3009_is_not_marked_applied(monkeypatch, tmp_path):
|
||||
_stub_v2_env(monkeypatch, tmp_path)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: "relation already exists")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_v2_roll_back_migration_best_effort",
|
||||
lambda name: pytest.fail("duplicate-object errors do not prove rollback is safe"),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_resolve_specific_migration",
|
||||
lambda name: pytest.fail("duplicate-object errors do not prove all SQL completed"),
|
||||
)
|
||||
_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)
|
||||
|
|
@ -197,28 +179,15 @@ def test_v2_duplicate_object_p3009_is_not_marked_applied(monkeypatch, tmp_path):
|
|||
|
||||
|
||||
def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path):
|
||||
"""v2 must never call _resolve_all_migrations — that's the bug it fixes."""
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
_stub_v2_env(monkeypatch, tmp_path)
|
||||
run = Mock(side_effect=_succeed_after(0, ""))
|
||||
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", run)
|
||||
|
||||
class FakeResult:
|
||||
stdout = "Applied migration.\n"
|
||||
stderr = ""
|
||||
|
||||
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", lambda *a, **kw: FakeResult())
|
||||
|
||||
resolve_called = {"n": 0}
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_resolve_all_migrations",
|
||||
lambda *a, **kw: resolve_called.__setitem__("n", resolve_called["n"] + 1),
|
||||
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"],
|
||||
)
|
||||
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
assert ok is True
|
||||
assert resolve_called["n"] == 0, "v2 must not invoke the diff-and-force recovery"
|
||||
|
||||
|
||||
_DEADLOCK_P3018_STDERR = (
|
||||
"Error: P3018\n"
|
||||
|
|
@ -228,12 +197,34 @@ _DEADLOCK_P3018_STDERR = (
|
|||
)
|
||||
|
||||
|
||||
def _stub_v2_env(monkeypatch, tmp_path):
|
||||
def _stub_v2_env(monkeypatch, tmp_path, ledger_logs=""):
|
||||
import psycopg
|
||||
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
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):
|
||||
|
|
@ -259,28 +250,21 @@ def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path):
|
|||
instance rolls the ledger row back and retries instead of dying."""
|
||||
_stub_v2_env(monkeypatch, tmp_path)
|
||||
|
||||
rolled_back = []
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_v2_roll_back_migration_best_effort",
|
||||
lambda name: rolled_back.append(name),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_resolve_specific_migration",
|
||||
lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
|
||||
)
|
||||
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, _DEADLOCK_P3018_STDERR))
|
||||
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 rolled_back == ["20260415120000_health_check_latest_per_model_index"]
|
||||
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)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_v2_roll_back_migration_best_effort", lambda name: None)
|
||||
|
||||
with patch(
|
||||
"litellm_proxy_extras.prisma_toolchain.run_prisma",
|
||||
|
|
@ -293,7 +277,7 @@ def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path):
|
|||
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)
|
||||
_stub_v2_env(monkeypatch, tmp_path, ledger_logs="ERROR: deadlock detected\nDETAIL: Process 72 waits for ShareLock")
|
||||
|
||||
stderr = (
|
||||
"Error: P3009\n"
|
||||
|
|
@ -301,27 +285,16 @@ def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_
|
|||
"The `20260415120000_health_check_latest_per_model_index` migration "
|
||||
"started at 2026-09-01 18:46:13 UTC failed"
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_failed_migration_logs",
|
||||
lambda name: "ERROR: deadlock detected\nDETAIL: Process 72 waits for ShareLock",
|
||||
)
|
||||
rolled_back = []
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_v2_roll_back_migration_best_effort",
|
||||
lambda name: rolled_back.append(name),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_resolve_specific_migration",
|
||||
lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
|
||||
)
|
||||
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr))
|
||||
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 rolled_back == ["20260415120000_health_check_latest_per_model_index"]
|
||||
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):
|
||||
|
|
@ -332,12 +305,6 @@ def test_v2_p3009_empty_ledger_logs_do_not_prove_completion(monkeypatch, tmp_pat
|
|||
"The `20260415120000_health_check_latest_per_model_index` migration "
|
||||
"started at 2026-09-01 18:46:13 UTC failed"
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: "")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_v2_roll_back_migration_best_effort",
|
||||
lambda name: pytest.fail("empty logs do not prove rollback is safe"),
|
||||
)
|
||||
with patch(
|
||||
"litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)
|
||||
) as run:
|
||||
|
|
@ -350,7 +317,7 @@ def test_v2_p3009_empty_ledger_logs_do_not_prove_completion(monkeypatch, tmp_pat
|
|||
|
||||
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)
|
||||
_stub_v2_env(monkeypatch, tmp_path, ledger_logs=None)
|
||||
|
||||
stderr = (
|
||||
"Error: P3009\n"
|
||||
|
|
@ -358,12 +325,6 @@ def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path):
|
|||
"The `20260415120000_health_check_latest_per_model_index` migration "
|
||||
"started at 2026-09-01 18:46:13 UTC failed"
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: None)
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_v2_roll_back_migration_best_effort",
|
||||
lambda name: pytest.fail("an unreadable ledger must not trigger a retry"),
|
||||
)
|
||||
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr))
|
||||
|
||||
with pytest.raises(RuntimeError, match="Migration completion could not be verified"):
|
||||
|
|
@ -372,7 +333,7 @@ def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path):
|
|||
|
||||
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)
|
||||
_stub_v2_env(monkeypatch, tmp_path, ledger_logs='ERROR: syntax error at or near "BRKN"')
|
||||
|
||||
stderr = (
|
||||
"Error: P3009\n"
|
||||
|
|
@ -380,11 +341,6 @@ def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path):
|
|||
"The `20260101000000_genuinely_broken` migration started at "
|
||||
"2026-09-01 18:46:13 UTC failed"
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_failed_migration_logs",
|
||||
lambda name: 'ERROR: syntax error at or near "BRKN"',
|
||||
)
|
||||
|
||||
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"):
|
||||
|
|
|
|||
|
|
@ -112,6 +112,7 @@ def select_tests(changed: tuple[str, ...]) -> tuple[str, ...]:
|
|||
(
|
||||
(("tests/e2e/logging/test_datadog_e2e.py", "litellm/router.py"), ("tests/e2e/logging/test_datadog_e2e.py",)),
|
||||
(("tests/e2e/ui/test_keys.py", "tests/e2e/claude_code/test_cli.py", "tests/e2e/load/test_burst.py"), ()),
|
||||
(("tests/e2e/migrations/test_startup.py", "tests/e2e/migrations/test_recovery.py"), ()),
|
||||
(("tests/e2e/batches/test_managed_files_enforcement_e2e.py",), ()),
|
||||
(("tests/e2e/guardrails/test_presidio_masking_e2e.py",), ()),
|
||||
(("tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py",), ()),
|
||||
|
|
@ -157,6 +158,10 @@ def test_a_changed_canary_file_is_selected_once_alongside_a_harness_change() ->
|
|||
assert select_tests((CANARY[1], "tests/e2e/proxy_client.py")) == CANARY
|
||||
|
||||
|
||||
def test_dedicated_migration_tests_do_not_suppress_shared_harness_canaries() -> None:
|
||||
assert select_tests(("tests/e2e/migrations/test_startup.py", "tests/e2e/conftest.py")) == CANARY
|
||||
|
||||
|
||||
def test_the_canary_joins_directly_selected_files_in_sorted_order() -> None:
|
||||
assert select_tests(("tests/e2e/logging/test_datadog_e2e.py", ".github/e2e-stack/up.sh")) == (
|
||||
*CANARY,
|
||||
|
|
|
|||
|
|
@ -64,7 +64,9 @@ def jwt_identity(idp: Keycloak, resources: ResourceManager, proxy: ProxyClient)
|
|||
|
||||
|
||||
def pytest_configure(config: pytest.Config) -> None:
|
||||
config.addinivalue_line("markers", "migration_startup: isolated container startup tests run by the migration CI workflow")
|
||||
config.addinivalue_line(
|
||||
"markers", "migration_startup: isolated container startup tests run by the migration CI workflow"
|
||||
)
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"e2e: live test that requires a running proxy and real provider keys",
|
||||
|
|
|
|||
|
|
@ -29,7 +29,8 @@ def start_replicas(
|
|||
|
||||
def assert_completed(database: Database, migration: Migration = COMPLETE) -> None:
|
||||
assert database.query(
|
||||
"SELECT finished_at IS NOT NULL, rolled_back_at IS NULL, applied_steps_count FROM _prisma_migrations WHERE migration_name = %s",
|
||||
'SELECT finished_at IS NOT NULL, rolled_back_at IS NULL, applied_steps_count FROM '
|
||||
'_prisma_migrations WHERE migration_name = %s',
|
||||
(migration.name,),
|
||||
) == ((True, True, 1),), "Expected exactly one successful SQL execution"
|
||||
assert database.query("SELECT id FROM migration_effect") == ((1,),)
|
||||
|
|
@ -47,7 +48,8 @@ def confirmed_history(database: Database) -> str:
|
|||
|
||||
def assert_original_proof(database: Database, row_id: str, finished: bool) -> None:
|
||||
assert database.query(
|
||||
"SELECT id, applied_steps_count, finished_at IS NOT NULL, rolled_back_at IS NULL FROM _prisma_migrations WHERE migration_name = %s",
|
||||
'SELECT id, applied_steps_count, finished_at IS NOT NULL, rolled_back_at IS NULL FROM '
|
||||
'_prisma_migrations WHERE migration_name = %s',
|
||||
(COMPLETE.name,),
|
||||
) == ((row_id, 1, finished, True),), "Recovery lost or replaced the original durable SQL proof"
|
||||
assert database.query("SELECT id FROM migration_effect") == ((1,),)
|
||||
|
|
@ -59,7 +61,8 @@ def pause_completion(database: Database) -> None:
|
|||
"CREATE FUNCTION migration_pause() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN "
|
||||
"IF NEW.migration_name = {name} AND NEW.finished_at IS NOT NULL THEN "
|
||||
"PERFORM pg_advisory_lock({gate}); PERFORM pg_advisory_unlock({gate}); END IF; RETURN NEW; END $$; "
|
||||
"CREATE TRIGGER migration_pause BEFORE UPDATE ON _prisma_migrations FOR EACH ROW EXECUTE FUNCTION migration_pause()"
|
||||
'CREATE TRIGGER migration_pause BEFORE UPDATE ON _prisma_migrations FOR EACH ROW '
|
||||
'EXECUTE FUNCTION migration_pause()'
|
||||
).format(name=sql.Literal(COMPLETE.name), gate=sql.Literal(GATE_KEY))
|
||||
)
|
||||
|
||||
|
|
@ -105,7 +108,8 @@ def unconfirmed(replicas: tuple[Replica, ...], database: Database) -> None:
|
|||
failed(replicas, "Migration completion could not be verified")
|
||||
started: Final = str(
|
||||
database.query(
|
||||
"SELECT to_char(started_at AT TIME ZONE 'UTC', 'YYYY-MM-DD HH24:MI:SS') FROM _prisma_migrations WHERE migration_name = %s",
|
||||
"SELECT to_char(started_at AT TIME ZONE 'UTC', 'YYYY-MM-DD HH24:MI:SS') FROM "
|
||||
'_prisma_migrations WHERE migration_name = %s',
|
||||
(COMPLETE.name,),
|
||||
)[0][0]
|
||||
)
|
||||
|
|
|
|||
|
|
@ -15,7 +15,9 @@ def adopt_legacy(containers: Containers, database: Database) -> None:
|
|||
count: Final = database.query("SELECT count(*) FROM _prisma_migrations")[0][0]
|
||||
existing_keys: Final = database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token')
|
||||
database.execute(
|
||||
"INSERT INTO \"LiteLLM_ShadowEvalJob\" (id, group_id, target_id, router_name, judge_model, shadow_percentage, max_turns, ends_at, stopped_at) VALUES ('migration-legacy', 'migration-legacy', 'target', 'router', 'judge', 1, 1, now(), now())"
|
||||
'INSERT INTO "LiteLLM_ShadowEvalJob" (id, group_id, target_id, router_name, judge_model, '
|
||||
"shadow_percentage, max_turns, ends_at, stopped_at) VALUES ('migration-legacy', "
|
||||
"'migration-legacy', 'target', 'router', 'judge', 1, 1, now(), now())"
|
||||
)
|
||||
database.execute("DROP TABLE _prisma_migrations")
|
||||
with ExitStack() as stack:
|
||||
|
|
@ -30,7 +32,8 @@ def adopt_legacy(containers: Containers, database: Database) -> None:
|
|||
assert detail in logs
|
||||
assert database.query("SELECT count(*) FROM _prisma_migrations") == ((count,),)
|
||||
assert database.query(
|
||||
"SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NULL OR rolled_back_at IS NOT NULL OR applied_steps_count <> 0"
|
||||
'SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NULL OR rolled_back_at IS '
|
||||
'NOT NULL OR applied_steps_count <> 0'
|
||||
) == ((0,),)
|
||||
assert set(existing_keys).issubset(database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token'))
|
||||
assert database.query("SELECT stopped_by FROM \"LiteLLM_ShadowEvalJob\" WHERE id = 'migration-legacy'") == (
|
||||
|
|
|
|||
|
|
@ -50,7 +50,8 @@ def pool(database: Database, output: Path) -> Generator[str]:
|
|||
f"[databases]\n* = host={url.hostname} port={url.port} user={url.username} password={url.password}\n"
|
||||
"[pgbouncer]\nlisten_addr = 0.0.0.0\nlisten_port = 6432\nauth_type = trust\nauth_file = /pool/users.txt\n"
|
||||
"pool_mode = transaction\ndefault_pool_size = 1\nreserve_pool_size = 0\nmax_client_conn = 100\n"
|
||||
"max_prepared_statements = 100\nquery_wait_timeout = 8\nignore_startup_parameters = extra_float_digits,options\n"
|
||||
'max_prepared_statements = 100\nquery_wait_timeout = 8\nignore_startup_parameters = '
|
||||
'extra_float_digits,options\n'
|
||||
)
|
||||
try:
|
||||
docker(
|
||||
|
|
|
|||
|
|
@ -159,7 +159,9 @@ class TestMigrationRecovery:
|
|||
)
|
||||
case "duplicate_history":
|
||||
database.execute(
|
||||
"INSERT INTO _prisma_migrations (id, migration_name, checksum, applied_steps_count) SELECT %s, migration_name, checksum, applied_steps_count FROM _prisma_migrations WHERE migration_name = %s",
|
||||
'INSERT INTO _prisma_migrations (id, migration_name, checksum, '
|
||||
'applied_steps_count) SELECT %s, migration_name, checksum, '
|
||||
'applied_steps_count FROM _prisma_migrations WHERE migration_name = %s',
|
||||
(str(uuid4()), COMPLETE.name),
|
||||
)
|
||||
case "missing_script":
|
||||
|
|
|
|||
|
|
@ -56,7 +56,8 @@ class TestMigrationStartup:
|
|||
replicas: Final = start_replicas(stack, containers, database, (FATAL,))
|
||||
failed(replicas, COMPLETE.name)
|
||||
assert database.query(
|
||||
"SELECT count(*) FROM _prisma_migrations WHERE migration_name = %s AND logs LIKE %s AND finished_at IS NULL",
|
||||
'SELECT count(*) FROM _prisma_migrations WHERE migration_name = %s AND logs LIKE '
|
||||
'%s AND finished_at IS NULL',
|
||||
(COMPLETE.name, "%MIGRATION_TEST_FATAL%"),
|
||||
) == ((1,),)
|
||||
|
||||
|
|
|
|||
|
|
@ -771,17 +771,7 @@ class _MigrateDeployHarness:
|
|||
self.confirmed_migrations = set(confirmed_migrations)
|
||||
|
||||
monkeypatch.delenv("DATABASE_URL", raising=False)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", staticmethod(lambda: str(tmp_path)))
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_roll_back_migration",
|
||||
staticmethod(lambda name: None),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_resolve_specific_migration",
|
||||
staticmethod(self.resolved.append),
|
||||
)
|
||||
monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(tmp_path))
|
||||
monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", self._fake_run)
|
||||
monkeypatch.setattr(utils_module, "_get_prisma_env", lambda: {})
|
||||
monkeypatch.setattr(utils_module.time, "sleep", lambda seconds: None)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue