litellm/tests/e2e/migrations/test_legacy.py
ryan-crabbe-berri 40a9b959a6
test(e2e): tag a2a, access_control, other, secret_manager and migrations tests with Subject metadata (#44965)
* test(e2e): add enum values, auto-discovering label gates and secret hiding for e2e metadata

* test(e2e): tag a2a, access_control, other, secret_manager and migrations tests with Subject metadata

* docs(e2e): name every markerless harness test file that carries no Subject

* test(e2e): keep the step discovery comprehensions to one for clause
2026-10-07 12:11:55 -07:00

103 lines
4.6 KiB
Python

from contextlib import ExitStack
from dataclasses import replace
from typing import Final, Literal
import pytest
from .checks import COMPLETE, assert_completed, confirmed_history, assert_original_proof, start_replicas
from .containers import Containers, failed, ready, seeded
from .database import Database, Databases
from e2e_metadata import Domain, Subject, meta
pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]
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())"
)
database.execute("DROP TABLE _prisma_migrations")
with ExitStack() as stack:
replicas: Final = start_replicas(stack, containers, database)
ready(replicas, database)
logs: Final = "\n".join(replica.logs() for replica in replicas)
for detail in (
"Legacy migration history was missing",
"historical data backfills were not replayed or verified",
"Continuing startup",
):
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'
) == ((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'") == (
(None,),
)
class TestLegacyMigrations:
@meta(
Subject(
domain=Domain.DB,
)
)
def test_matching_schema_warns_and_starts(self, containers: Containers, database: Database) -> None:
adopt_legacy(containers, database)
@pytest.mark.parametrize("fault", ("schema_drift", "custom_migrations", "empty_ledger"))
@meta(
Subject(
domain=Domain.DB,
)
)
def test_unrecognized_legacy_state_is_not_baselined(
self, containers: Containers, database: Database, fault: str
) -> None:
if fault == "empty_ledger":
database.execute("TRUNCATE _prisma_migrations")
else:
database.execute("DROP TABLE _prisma_migrations")
if fault == "schema_drift":
database.execute('ALTER TABLE "LiteLLM_VerificationToken" DROP COLUMN key_alias CASCADE')
with containers.start(database, (COMPLETE,) if fault == "custom_migrations" else ()) as replica:
failed((replica,), "Cannot automatically baseline" if fault != "empty_ledger" else "migration")
assert not database.exists("migration_effect")
if database.exists("_prisma_migrations"):
assert database.query(
"SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NOT NULL AND applied_steps_count <> 1"
) == ((0,),)
@pytest.mark.parametrize("scenario", ("upgrade", "recovery", "legacy"))
@meta(
Subject(
domain=Domain.DB,
)
)
def test_non_default_schema(
self, containers: Containers, databases: Databases, scenario: Literal["upgrade", "recovery", "legacy"]
) -> None:
with databases.create(schema="migration tenant") as database:
with containers.start(database) as seed:
seeded(seed, database)
match scenario:
case "upgrade":
with ExitStack() as stack:
ready(start_replicas(stack, containers, database, (COMPLETE,)), database)
assert_completed(database)
case "recovery":
original: Final = confirmed_history(database)
with ExitStack() as stack:
ready(start_replicas(stack, containers, database, (COMPLETE,)), database)
assert_original_proof(database, original, True)
case "legacy":
adopt_legacy(containers, database)
public: Final = replace(database, schema="public")
assert not public.exists("_prisma_migrations")
assert not public.exists('"LiteLLM_VerificationToken"')