litellm/tests/e2e/migrations/test_shaped_database.py
Yuneng Jiang 45d22dc5e1
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.
2026-09-21 12:50:58 -07:00

43 lines
1.8 KiB
Python

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"