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 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"


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()


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

    def state(self) -> ContainerState:
        return ContainerState.model_validate_json(docker("inspect", "--format", "{{json .State}}", self.name))

    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)

    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

    def kill(self) -> None:
        if self.state().Running:
            docker("kill", self.name)

    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,),)


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)


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}"


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)

    @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)