mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
45d22dc5e1
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.
91 lines
3.9 KiB
Python
91 lines
3.9 KiB
Python
import json
|
|
import os
|
|
from collections.abc import Iterator
|
|
from pathlib import Path
|
|
from typing import Final
|
|
from urllib.parse import urlsplit
|
|
|
|
import pytest
|
|
from _pytest.fixtures import SubRequest
|
|
|
|
from .containers import Containers, docker, ready
|
|
from .database import Database, Databases
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def migration_image(tmp_path_factory: pytest.TempPathFactory) -> str:
|
|
configured: Final = os.environ.get("LITELLM_MIGRATION_TEST_IMAGE")
|
|
assert configured, "LITELLM_MIGRATION_TEST_IMAGE must name the built candidate image"
|
|
image: Final = docker("image", "inspect", configured, "--format", "{{.Id}}")
|
|
assert image.startswith("sha256:"), "Unable to identify the candidate image"
|
|
output: Final = Path(os.environ.get("MIGRATION_TEST_OUTPUT", str(tmp_path_factory.getbasetemp())))
|
|
output.mkdir(parents=True, exist_ok=True)
|
|
(output / "image.json").write_text(json.dumps({"requested": configured, "image_id": image}))
|
|
return image
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def databases() -> Databases:
|
|
admin: Final = os.environ.get("MIGRATION_TEST_ADMIN_URL", "")
|
|
parsed: Final = urlsplit(admin)
|
|
assert parsed.hostname in ("127.0.0.1", "localhost"), "Use an isolated loopback PostgreSQL test cluster"
|
|
assert parsed.port and parsed.path and not parsed.query, "Supply the test cluster port and admin database"
|
|
container_admin: Final = os.environ.get(
|
|
"MIGRATION_TEST_CONTAINER_ADMIN_URL",
|
|
admin.replace("127.0.0.1", "host.docker.internal").replace("localhost", "host.docker.internal"),
|
|
)
|
|
return Databases(admin, container_admin)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def migrated_template(
|
|
databases: Databases, migration_image: str, tmp_path_factory: pytest.TempPathFactory
|
|
) -> Iterator[Database]:
|
|
output: Final = Path(os.environ.get("MIGRATION_TEST_OUTPUT", str(tmp_path_factory.getbasetemp()))) / "seed"
|
|
with databases.create() as database:
|
|
with Containers(migration_image, output).start(database) as replica:
|
|
ready((replica,), database)
|
|
yield database
|
|
|
|
|
|
@pytest.fixture
|
|
def database(databases: Databases, migrated_template: Database) -> Iterator[Database]:
|
|
with databases.create(migrated_template) as database:
|
|
yield database
|
|
|
|
|
|
@pytest.fixture
|
|
def containers(migration_image: str, tmp_path: Path, request: SubRequest) -> Containers:
|
|
configured: Final = os.environ.get("MIGRATION_TEST_OUTPUT")
|
|
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
|