diff --git a/.github/e2e-stack/select_tests.py b/.github/e2e-stack/select_tests.py index 238818a0d36..9dd880c05cd 100644 --- a/.github/e2e-stack/select_tests.py +++ b/.github/e2e-stack/select_tests.py @@ -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$" diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 2749db5d754..8a83c786e02 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.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) diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py index 338c571eb4f..832075f6fbe 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -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"): diff --git a/tests/code_coverage_tests/test_e2e_changed_gate.py b/tests/code_coverage_tests/test_e2e_changed_gate.py index 588402e3996..a628fbb0633 100644 --- a/tests/code_coverage_tests/test_e2e_changed_gate.py +++ b/tests/code_coverage_tests/test_e2e_changed_gate.py @@ -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, diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 4a5f0aa880f..53d35effdc9 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -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", diff --git a/tests/e2e/migrations/checks.py b/tests/e2e/migrations/checks.py index b14631ccad8..619ad3b0e6c 100644 --- a/tests/e2e/migrations/checks.py +++ b/tests/e2e/migrations/checks.py @@ -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] ) diff --git a/tests/e2e/migrations/test_legacy.py b/tests/e2e/migrations/test_legacy.py index 7ba73eb82e0..ba5e77a3070 100644 --- a/tests/e2e/migrations/test_legacy.py +++ b/tests/e2e/migrations/test_legacy.py @@ -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'") == ( diff --git a/tests/e2e/migrations/test_pooling.py b/tests/e2e/migrations/test_pooling.py index e0a3693b33e..c4015549de3 100644 --- a/tests/e2e/migrations/test_pooling.py +++ b/tests/e2e/migrations/test_pooling.py @@ -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( diff --git a/tests/e2e/migrations/test_recovery.py b/tests/e2e/migrations/test_recovery.py index 58bf5c348d6..80e5747eaac 100644 --- a/tests/e2e/migrations/test_recovery.py +++ b/tests/e2e/migrations/test_recovery.py @@ -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": diff --git a/tests/e2e/migrations/test_startup.py b/tests/e2e/migrations/test_startup.py index a648218cb26..dc628a8ed7a 100644 --- a/tests/e2e/migrations/test_startup.py +++ b/tests/e2e/migrations/test_startup.py @@ -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,),) diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 57133ea95c4..bb329264a11 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -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)