test(migrations): cover the release-to-release upgrade path

The migration e2e harness only ever used one image: it seeded the database
with the candidate build and then applied synthetic migrations on top. That
proves the migration machinery (locking, crash recovery, legacy baselining,
pooling) but never executes the real schema of release N against the real
migrations of release N+1, which is the path operators actually run.

Adds a baseline image alongside the candidate, so a test can seed with a
published release and upgrade with the build under test.

Suites:

- test_upgrade.py: the candidate applies the pending release migrations,
  keys minted by the baseline release survive, and concurrent replicas
  upgrade a baseline database exactly once.
- test_rolling_upgrade.py: a baseline replica keeps serving virtual-key
  auth while the candidate migrates underneath it, and both releases serve
  and resolve each other's keys during the overlap. This is the reported
  failure: a new column on LiteLLM_VerificationToken invalidates prepared
  plans on pods still running the old release, which the proxy reads
  whole-row, and auth starts failing until those pods leave service.
- test_shaped_database.py: the upgrade completes and preserves rows on a
  populated spend log, rather than on the empty database every other
  migration test starts from.

Every upgrade assertion is gated on the candidate having actually applied
migrations the baseline had not, so a stale pin fails loudly instead of
passing on an empty delta.

CI adds two jobs to the migration_startup workflow. The baseline defaults
to a committed release pin and is overridable per pipeline, matching how
migration_candidate_image already works; only the upgrade jobs pull it.

Verified against a real v1.101.0 -> v1.102.0 upgrade: 6 passed, with the
baseline seeding 165 migrations and the candidate applying the 6 that
landed between the two releases.

.circleci/config.yml

@@ -6,6 +6,9 @@ parameters:
   migration_candidate_image:
     type: string
     default: ""
+  migration_baseline_image:
+    type: string
+    default: "ghcr.io/berriai/litellm-database:v1.102.0"
   migration_source_sha:
     type: string
     default: ""
@@ -2946,7 +2949,10 @@ jobs:
     parameters:
       suite:
         type: enum
-        enum: [startup, recovery, legacy]
+        enum: [startup, recovery, legacy, upgrade, shaped]
+      baseline:
+        type: boolean
+        default: false
     machine:
       image: ubuntu-2204:2024.04.1
     resource_class: large
@@ -2954,6 +2960,7 @@ jobs:
     environment:
       LITELLM_MIGRATION_TESTS: "1"
       LITELLM_MIGRATION_TEST_IMAGE: litellm-docker-database:ci
+      LITELLM_MIGRATION_BASELINE_IMAGE: << pipeline.parameters.migration_baseline_image >>
       MIGRATION_TEST_ADMIN_URL: postgresql://postgres:postgres@127.0.0.1:5432/postgres
       MIGRATION_TEST_CONTAINER_ADMIN_URL: postgresql://postgres:postgres@host.docker.internal:5432/postgres
       MIGRATION_TEST_OUTPUT: /tmp/migration-results
@@ -2981,6 +2988,16 @@ jobs:
       - wait_for_service:
           url: tcp://localhost:5432
           timeout: "60"
+      - when:
+          condition: << parameters.baseline >>
+          steps:
+            - run:
+                name: Pull the baseline release the upgrade starts from
+                environment:
+                  BASELINE_IMAGE: << pipeline.parameters.migration_baseline_image >>
+                command: |
+                  [[ "$BASELINE_IMAGE" =~ ^ghcr.io/berriai/[a-z0-9._/-]+(@sha256:[0-9a-f]{64}|:v[0-9][0-9a-z.-]*)$ ]] || exit 1
+                  docker pull "$BASELINE_IMAGE"
       - run:
           name: Run migration startup regressions
           environment:
@@ -3188,6 +3205,16 @@ workflows:
           name: migration-legacy-and-pooling
           suite: legacy
           requires: [build_docker_database_image]
+      - migration_startup_tests:
+          name: migration-upgrade
+          suite: upgrade
+          baseline: true
+          requires: [build_docker_database_image]
+      - migration_startup_tests:
+          name: migration-upgrade-shaped
+          suite: shaped
+          baseline: true
+          requires: [build_docker_database_image]
   migration_startup_scheduled:
     triggers:
       - schedule:

.circleci/scripts/run_migration_tests.py

@@ -13,6 +13,8 @@ SUITES: Final = {
     "startup": (("test_startup.py",), 12),
     "recovery": (("test_recovery.py",), 15),
     "legacy": (("test_legacy.py", "test_pooling.py"), 11),
+    "upgrade": (("test_upgrade.py", "test_rolling_upgrade.py"), 5),
+    "shaped": (("test_shaped_database.py",), 1),
 }
 
 
@@ -93,6 +95,7 @@ def main() -> int:
             {
                 **metadata,
                 "suite": suite,
+                "baseline_image": os.environ.get("LITELLM_MIGRATION_BASELINE_IMAGE", ""),
                 "expected_cases": expected,
                 "passed": passed,
                 "pytest_exit_code": result.returncode,

tests/e2e/migrations/conftest.py

@@ -60,3 +60,32 @@ def containers(migration_image: str, tmp_path: Path, request: SubRequest) -> Con
     output: Final = Path(configured) / request.node.name if configured else tmp_path
     output.mkdir(parents=True, exist_ok=True)
     return Containers(migration_image, output)
+
+
+@pytest.fixture(scope="session")
+def baseline_image(tmp_path_factory: pytest.TempPathFactory) -> str:
+    configured: Final = os.environ.get("LITELLM_MIGRATION_BASELINE_IMAGE")
+    assert configured, "LITELLM_MIGRATION_BASELINE_IMAGE must name the released image the upgrade starts from"
+    image: Final = docker("image", "inspect", configured, "--format", "{{.Id}}")
+    assert image.startswith("sha256:"), "Unable to identify the baseline image"
+    output: Final = Path(os.environ.get("MIGRATION_TEST_OUTPUT", str(tmp_path_factory.getbasetemp())))
+    output.mkdir(parents=True, exist_ok=True)
+    (output / "baseline-image.json").write_text(json.dumps({"requested": configured, "image_id": image}))
+    return image
+
+
+@pytest.fixture(scope="session")
+def baseline_template(
+    databases: Databases, baseline_image: str, tmp_path_factory: pytest.TempPathFactory
+) -> Iterator[Database]:
+    output: Final = Path(os.environ.get("MIGRATION_TEST_OUTPUT", str(tmp_path_factory.getbasetemp()))) / "baseline-seed"
+    with databases.create() as database:
+        with Containers(baseline_image, output).start(database) as replica:
+            ready((replica,), database)
+        yield database
+
+
+@pytest.fixture
+def baseline_database(databases: Databases, baseline_template: Database) -> Iterator[Database]:
+    with databases.create(baseline_template) as database:
+        yield database

tests/e2e/migrations/containers.py

@@ -5,7 +5,7 @@ import subprocess
 import time
 from collections.abc import Callable, Generator, Mapping
 from contextlib import contextmanager
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from pathlib import Path
 from typing import Final
 from uuid import uuid4
@@ -123,6 +123,9 @@ class Containers:
     image: str
     output: Path
 
+    def using(self, image: str) -> "Containers":
+        return replace(self, image=image)
+
     @contextmanager
     def start(
         self,

tests/e2e/migrations/test_rolling_upgrade.py

@@ -0,0 +1,55 @@
+from typing import Final
+
+import pytest
+
+from .containers import Containers, ready
+from .database import Database
+from .upgrade import (
+    CACHED_PLAN,
+    assert_history_clean,
+    assert_upgraded,
+    auth_traffic,
+    confirm,
+    keep_serving,
+    migration_names,
+    provision,
+)
+
+pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]
+
+
+class TestRollingUpgrade:
+    def test_baseline_replica_keeps_serving_while_the_candidate_migrates(
+        self, containers: Containers, baseline_image: str, baseline_database: Database
+    ) -> None:
+        with containers.using(baseline_image).start(baseline_database) as old:
+            ready((old,), baseline_database)
+            key, _ = provision(old)
+            before: Final = migration_names(baseline_database)
+            with auth_traffic(old, key) as traffic:
+                keep_serving(traffic, "the baseline replica authenticating before the upgrade")
+                with containers.start(baseline_database) as new:
+                    ready((new,), baseline_database)
+                    assert_upgraded(before, migration_names(baseline_database))
+                    keep_serving(traffic, "the baseline replica authenticating after the schema moved")
+            assert_history_clean(baseline_database)
+            assert CACHED_PLAN not in old.logs(), "The baseline replica hit a stale prepared statement"
+            assert old.state().Running, "The baseline replica died during the upgrade"
+
+    def test_both_releases_serve_and_share_keys_during_the_overlap(
+        self, containers: Containers, baseline_image: str, baseline_database: Database
+    ) -> None:
+        with containers.using(baseline_image).start(baseline_database) as old:
+            ready((old,), baseline_database)
+            old_key, old_alias = provision(old)
+            before: Final = migration_names(baseline_database)
+            with containers.start(baseline_database) as new:
+                ready((new,), baseline_database)
+                assert_upgraded(before, migration_names(baseline_database))
+                new_key, new_alias = provision(new)
+                with auth_traffic(old, old_key) as old_traffic, auth_traffic(new, new_key) as new_traffic:
+                    keep_serving(old_traffic, "the baseline replica serving through the overlap")
+                    keep_serving(new_traffic, "the candidate replica serving through the overlap")
+                confirm(old, new_key, new_alias)
+                confirm(new, old_key, old_alias)
+            assert CACHED_PLAN not in old.logs(), "The baseline replica hit a stale prepared statement"

tests/e2e/migrations/test_shaped_database.py

@@ -0,0 +1,43 @@
+from typing import Final
+
+import pytest
+
+from .containers import Containers, ready
+from .database import Database
+from .upgrade import assert_history_clean, assert_upgraded, confirm, migration_names, provision
+
+SPEND_ROWS: Final = 20_000
+
+pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]
+
+
+def seed_spend_logs(database: Database, rows: int) -> None:
+    database.execute(
+        'INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, "startTime", "endTime") '
+        "SELECT 'upgrade-shape-' || g, 'acompletion', now() - (g || ' seconds')::interval, "
+        "now() - (g || ' seconds')::interval FROM generate_series(1, %s) AS g",
+        (rows,),
+    )
+    assert database.query('SELECT count(*) FROM "LiteLLM_SpendLogs"') == ((rows,),)
+
+
+class TestPopulatedDatabaseUpgrade:
+    def test_upgrade_completes_and_preserves_a_populated_spend_log(
+        self, containers: Containers, baseline_image: str, baseline_database: Database
+    ) -> None:
+        with containers.using(baseline_image).start(baseline_database) as old:
+            ready((old,), baseline_database)
+            key, alias = provision(old)
+        seed_spend_logs(baseline_database, SPEND_ROWS)
+        before: Final = migration_names(baseline_database)
+        with containers.start(baseline_database) as new:
+            ready((new,), baseline_database)
+            assert_upgraded(before, migration_names(baseline_database))
+            confirm(new, key, alias)
+        assert_history_clean(baseline_database)
+        assert baseline_database.query('SELECT count(*) FROM "LiteLLM_SpendLogs"') == ((SPEND_ROWS,),), (
+            "The upgrade lost spend rows"
+        )
+        assert baseline_database.query(
+            'SELECT count(*) FROM "LiteLLM_SpendLogs" WHERE "startTime" IS NULL OR "endTime" IS NULL'
+        ) == ((0,),), "The upgrade nulled timestamps on existing spend rows"

tests/e2e/migrations/test_upgrade.py

@@ -0,0 +1,47 @@
+from contextlib import ExitStack
+from typing import Final
+
+import pytest
+
+from .checks import start_replicas
+from .containers import Containers, ready
+from .database import Database
+from .upgrade import assert_history_clean, assert_upgraded, confirm, migration_names, provision
+
+pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]
+
+
+class TestReleaseUpgrade:
+    def test_candidate_applies_the_pending_release_migrations(
+        self, containers: Containers, baseline_database: Database
+    ) -> None:
+        before: Final = migration_names(baseline_database)
+        with containers.start(baseline_database) as replica:
+            ready((replica,), baseline_database)
+        assert_upgraded(before, migration_names(baseline_database))
+        assert_history_clean(baseline_database)
+
+    def test_upgrade_preserves_keys_minted_by_the_baseline_release(
+        self, containers: Containers, baseline_image: str, baseline_database: Database
+    ) -> None:
+        with containers.using(baseline_image).start(baseline_database) as old:
+            ready((old,), baseline_database)
+            key, alias = provision(old)
+            confirm(old, key, alias)
+        before: Final = migration_names(baseline_database)
+        with containers.start(baseline_database) as new:
+            ready((new,), baseline_database)
+            assert_upgraded(before, migration_names(baseline_database))
+            confirm(new, key, alias)
+
+    def test_concurrent_replicas_upgrade_a_baseline_database_once(
+        self, containers: Containers, baseline_database: Database
+    ) -> None:
+        before: Final = migration_names(baseline_database)
+        with ExitStack() as stack:
+            ready(start_replicas(stack, containers, baseline_database), baseline_database)
+        assert_upgraded(before, migration_names(baseline_database))
+        assert_history_clean(baseline_database)
+        assert baseline_database.query("SELECT count(*) FROM _prisma_migrations WHERE applied_steps_count > 1") == (
+            (0,),
+        ), "A migration was executed more than once across the upgrading replicas"

tests/e2e/migrations/upgrade.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import threading
+from collections.abc import Generator
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Final
+from uuid import uuid4
+
+from e2e_http import Result, Success, unwrap
+from models import (
+    KeyGenerateBody,
+    KeyGenerateResponse,
+    KeyInfoParams,
+    KeyInfoResponse,
+    ModelsListParams,
+    ModelsListResponse,
+)
+from pydantic import BaseModel
+
+from .containers import Replica, until
+from .database import Database
+
+CACHED_PLAN: Final = "cached plan must not change result type"
+
+
+def provision(replica: Replica) -> tuple[str, str]:
+    alias: Final = f"upgrade-{uuid4().hex}"
+    key: Final = unwrap(
+        replica.transport.post(
+            "/key/generate",
+            headers=replica.transport.master,
+            json=KeyGenerateBody(key_alias=alias),
+            response_type=KeyGenerateResponse,
+        )
+    ).key
+    return key, alias
+
+
+def confirm(replica: Replica, key: str, alias: str) -> None:
+    info: Final = unwrap(
+        replica.transport.get(
+            "/key/info",
+            headers=replica.transport.master,
+            params=KeyInfoParams(key=key),
+            response_type=KeyInfoResponse,
+        )
+    )
+    assert info.info.key_alias == alias, "Key minted on one release did not resolve on the other"
+
+
+@dataclass(slots=True)
+class Outcomes:
+    served: int = 0
+    failures: list[str] = field(default_factory=list)
+
+    def record(self, result: Result[BaseModel]) -> None:
+        match result:
+            case Success():
+                self.served += 1
+            case _:
+                self.failures.append(result.model_dump_json())
+
+
+@contextmanager
+def auth_traffic(replica: Replica, key: str, interval: float = 0.05) -> Generator[Outcomes]:
+    outcomes: Final = Outcomes()
+    stop: Final = threading.Event()
+
+    def drive() -> None:
+        while not stop.is_set():
+            outcomes.record(
+                replica.transport.get(
+                    "/v1/models",
+                    headers=replica.transport.bearer(key),
+                    params=ModelsListParams(),
+                    response_type=ModelsListResponse,
+                    timeout=10,
+                )
+            )
+            stop.wait(interval)
+
+    thread: Final = threading.Thread(target=drive, name="upgrade-auth-traffic", daemon=True)
+    thread.start()
+    try:
+        yield outcomes
+    finally:
+        stop.set()
+        thread.join(30)
+        assert not thread.is_alive(), "Auth traffic thread did not stop"
+
+
+def keep_serving(outcomes: Outcomes, description: str, calls: int = 20) -> int:
+    target: Final = outcomes.served + calls
+    until(description, lambda: outcomes.served >= target or bool(outcomes.failures))
+    assert not outcomes.failures, f"Virtual-key auth failed during {description}: {outcomes.failures[:5]}"
+    return outcomes.served
+
+
+def migration_names(database: Database) -> frozenset[str]:
+    return frozenset(str(row[0]) for row in database.query("SELECT migration_name FROM _prisma_migrations"))
+
+
+def assert_history_clean(database: Database) -> None:
+    assert database.query(
+        "SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NULL OR rolled_back_at IS NOT NULL"
+    ) == ((0,),), "The upgrade left an unfinished or rolled-back migration behind"
+
+
+def assert_upgraded(before: frozenset[str], after: frozenset[str]) -> frozenset[str]:
+    applied: Final = after - before
+    assert applied, (
+        "The candidate applied no migrations the baseline release had not: the pinned "
+        "LITELLM_MIGRATION_BASELINE_IMAGE is at or ahead of the candidate, so this suite proves nothing"
+    )
+    assert not before - after, "The upgrade removed migration history the baseline release had already applied"
+    return applied