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("poll until the condition holds")
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("inspect the replica container")
    def state(self) -> ContainerState:
        return ContainerState.model_validate_json(docker("inspect", "--format", "{{json .State}}", self.name))

    @step("probe the replica's state and 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 replica container 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 replica container")
    def kill(self) -> None:
        if self.state().Running:
            docker("kill", self.name)

    @step("confirm the replica mints a usable key")
    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("poll replicas for readiness")
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("poll replicas for a rejected startup")
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("poll replicas to confirm they keep waiting")
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 replica container")
    @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)


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)