litellm/tests/e2e/migrations/containers.py

from __future__ import annotations

import hashlib
import subprocess
import time
from collections.abc import Callable, Generator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Final
from uuid import uuid4

from e2e_http import NoBody, Success, unwrap
from e2e_metadata import step
from models import KeyGenerateBody, KeyGenerateResponse, KeyInfoParams, KeyInfoResponse
from transport import HttpTransport

from .database import Database, prisma_url
from .startup_models import ContainerState, Migration, Observation, Readiness

MASTER_KEY: Final = "sk-migration-ci-fixture"


@step("Run a docker command")
def docker(*args: str) -> str:
    result: Final = subprocess.run(("docker", *args), capture_output=True, text=True, timeout=90)
    assert result.returncode == 0, f"Docker operation failed: {result.stderr}"
    return result.stdout.strip()


@step("Wait for {description}")
def until(description: str, condition: Callable[[], bool], seconds: float = 150) -> None:
    deadline: Final = time.monotonic() + seconds
    while time.monotonic() < deadline:
        if condition():
            return
        time.sleep(0.25)
    raise AssertionError(f"Timed out waiting for {description}")


@dataclass(frozen=True, slots=True)
class Replica:
    name: str
    transport: HttpTransport
    output: Path

    @step("Read the proxy container's state from docker inspect")
    def state(self) -> ContainerState:
        return ContainerState.model_validate_json(docker("inspect", "--format", "{{json .State}}", self.name))

    @step("Check whether the proxy container is running and ready on /health/readiness")
    def observe(self) -> Observation:
        state: Final = self.state()
        result: Final = self.transport.get(
            "/health/readiness", headers=self.transport.master, params=NoBody(), response_type=Readiness, timeout=1
        )
        ready: Final = isinstance(result, Success) and result.data.status == "healthy" and result.data.db == "connected"
        return Observation(None if state.Running else state.ExitCode, ready)

    @step("Read the proxy container's logs")
    def logs(self) -> str:
        result: Final = subprocess.run(("docker", "logs", self.name), capture_output=True, text=True, timeout=30)
        assert result.returncode == 0, result.stderr
        return result.stdout + result.stderr

    @step("Kill the proxy container")
    def kill(self) -> None:
        if self.state().Running:
            docker("kill", self.name)

    @step("Generate a virtual key on the proxy container and read it back from /key/info and the database")
    def usable(self, database: Database) -> None:
        alias: Final = f"migration-{uuid4().hex}"
        key: Final = unwrap(
            self.transport.post(
                "/key/generate",
                headers=self.transport.master,
                json=KeyGenerateBody(key_alias=alias),
                response_type=KeyGenerateResponse,
            )
        ).key
        info: Final = unwrap(
            self.transport.get(
                "/key/info",
                headers=self.transport.master,
                params=KeyInfoParams(key=key),
                response_type=KeyInfoResponse,
            )
        )
        assert info.info.key_alias == alias
        assert database.query(
            'SELECT key_alias FROM "LiteLLM_VerificationToken" WHERE token = %s',
            (hashlib.sha256(key.encode()).hexdigest(),),
        ) == ((alias,),)


@step("Wait for every proxy container to be ready, then generate and read back a virtual key on each")
def ready(replicas: tuple[Replica, ...], database: Database) -> None:
    def all_ready() -> bool:
        observations: Final = tuple(replica.observe() for replica in replicas)
        assert all(item.exit_code is None for item in observations), "Replica exited before readiness"
        return all(item.ready for item in observations)

    until("every replica ready", all_ready)
    for replica in replicas:
        replica.usable(database)


@step("Wait for the seed proxy container to be ready and finish building its request-log indexes")
def seeded(seed: Replica, database: Database) -> None:
    ready((seed,), database)
    until("the seed replica to finish its request-log indexes", lambda: request_log_indexes_built(database))


def request_log_indexes_built(database: Database) -> bool:
    return database.query(
        "SELECT count(*) FROM pg_index x JOIN pg_class i ON i.oid = x.indexrelid "
        "JOIN pg_namespace n ON n.oid = i.relnamespace WHERE n.nspname = current_schema() AND x.indisvalid "
        "AND i.relname IN ('LiteLLM_SpendLogs_api_key_startTime_idx', 'LiteLLM_SpendLogs_litellm_call_id_idx')"
    ) == ((2,),)


@step('Wait for every proxy container to refuse to start, logging "{marker}"')
def failed(replicas: tuple[Replica, ...], marker: str) -> None:
    def all_stopped() -> bool:
        observations: Final = tuple(replica.observe() for replica in replicas)
        assert not any(item.ready for item in observations), "Failed migration exposed a ready proxy"
        return all(item.exit_code is not None for item in observations)

    until("every replica to reject startup", all_stopped)
    for replica in replicas:
        assert replica.state().ExitCode != 0, "Failed startup returned success"
        assert marker in replica.logs(), f"Startup failed outside the expected migration: {marker}"


@step("Check that every proxy container keeps waiting without serving for {seconds}s")
def waiting(replicas: tuple[Replica, ...], seconds: float) -> None:
    deadline: Final = time.monotonic() + seconds
    while time.monotonic() < deadline:
        assert all(item.exit_code is None and not item.ready for item in (replica.observe() for replica in replicas)), (
            "Contending replica exited or served early"
        )
        time.sleep(0.25)


@dataclass(frozen=True, slots=True)
class Containers:
    image: str
    output: Path

    def using(self, image: str) -> "Containers":
        return replace(self, image=image)

    @step("Start a proxy container on the test database")
    @contextmanager
    def start(
        self,
        database: Database,
        migrations: tuple[Migration, ...] = (),
        *,
        v2: bool = True,
        disabled: bool = False,
        environment: Mapping[str, str] | None = None,
    ) -> Generator[Replica]:
        name: Final = f"litellm-migration-{uuid4().hex[:16]}"
        directory: Final = self.output / name
        directory.mkdir(parents=True)
        for migration in migrations:
            write_migration(directory, migration)
        (directory / "config.yaml").write_text(
            "model_list: []\ngeneral_settings:\n  master_key: os.environ/LITELLM_MASTER_KEY\n"
        )
        env: Final = {
            "DATABASE_URL": prisma_url(database.container_url, database.schema),
            "LITELLM_MASTER_KEY": MASTER_KEY,
            "LITELLM_SALT_KEY": MASTER_KEY,
            "LITELLM_LOCAL_MODEL_COST_MAP": "True",
            "LITELLM_TELEMETRY": "False",
            "LITELLM_LOG": "INFO",
            "DATABASE_CONNECTION_POOL_LIMIT": "2",
            "DEFAULT_NUM_WORKERS_LITELLM_PROXY": "1",
            "USE_V2_MIGRATION_RESOLVER": str(v2).lower(),
            "DISABLE_SCHEMA_UPDATE": str(disabled).lower(),
            "LITELLM_MIGRATION_DIR": "/migration-test/prisma",
            "LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT": "180",
            **(environment or {}),
        }
        try:
            docker(
                "run",
                "-d",
                "--name",
                name,
                "--label",
                "litellm-migration-test=true",
                "--add-host",
                "host.docker.internal:host-gateway",
                "-p",
                "127.0.0.1::4000",
                "-v",
                f"{directory}:/migration-test",
                *(arg for key, value in env.items() for arg in ("-e", f"{key}={value}")),
                self.image,
                "--config",
                "/migration-test/config.yaml",
                "--host",
                "0.0.0.0",
                "--port",
                "4000",
            )
            port: Final = int(docker("port", name, "4000/tcp").rsplit(":", 1)[1])
            replica: Final = Replica(name, HttpTransport(f"http://127.0.0.1:{port}", MASTER_KEY, 15), directory)
            yield replica
        finally:
            try:
                state: Final = subprocess.run(
                    ("docker", "inspect", "--format", "{{json .State}}", name),
                    capture_output=True,
                    text=True,
                    timeout=30,
                )
                (directory / "state.json").write_text(state.stdout or state.stderr)
                logs: Final = subprocess.run(("docker", "logs", name), capture_output=True, text=True, timeout=30)
                (directory / "proxy.log").write_text(logs.stdout + logs.stderr)
            finally:
                subprocess.run(("docker", "rm", "-f", name), capture_output=True, text=True, timeout=30, check=True)


@step("Write the migration {migration.name} into the proxy container's migration directory")
def write_migration(directory: Path, migration: Migration) -> None:
    path: Final = directory / "prisma" / "migrations" / migration.name
    path.mkdir(parents=True)
    (path / "migration.sql").write_text(migration.script)