scripts/docker_smoke.py

"""Build and smoke-test the Docker image in a cross-platform way."""

from __future__ import annotations

import argparse
import subprocess
import sys
import time
import urllib.request
from collections.abc import Sequence


def run(command: Sequence[str], *, check: bool = True) -> subprocess.CompletedProcess[str]:
    return subprocess.run(command, check=check, text=True)


def output(command: Sequence[str]) -> str:
    return subprocess.check_output(command, text=True).strip()


def print_logs(container_name: str) -> None:
    subprocess.run(["docker", "logs", container_name], check=False)


def remove_container(container_name: str) -> None:
    subprocess.run(
        ["docker", "rm", "-f", container_name],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        check=False,
    )


def wait_for_health(container_name: str, health_url: str, attempts: int, delay: float) -> int:
    for attempt in range(1, attempts + 1):
        status = output(["docker", "inspect", "-f", "{{.State.Status}}", container_name])
        if status != "running":
            print(f"Container stopped before becoming healthy. Status: {status}", file=sys.stderr)
            print_logs(container_name)
            return 1

        try:
            with urllib.request.urlopen(health_url, timeout=3) as response:
                if 200 <= response.status < 300:
                    print(f"Docker health endpoint passed on attempt {attempt}.")
                    subprocess.run(["docker", "logs", "--tail=80", container_name], check=False)
                    return 0
        except Exception as exc:
            print(f"Waiting for health endpoint, attempt {attempt}/{attempts}: {exc}")
            time.sleep(delay)

    print("Timed out waiting for Streamlit health endpoint.", file=sys.stderr)
    subprocess.run(["docker", "inspect", container_name], check=False)
    print_logs(container_name)
    return 1


def main() -> int:
    parser = argparse.ArgumentParser(description="Build and smoke-test the Docker image.")
    parser.add_argument("--image", default="rag-qa-command-center:ci")
    parser.add_argument("--port", default="8501")
    parser.add_argument("--attempts", type=int, default=45)
    parser.add_argument("--delay", type=float, default=2.0)
    args = parser.parse_args()

    container_name = f"rag-qa-command-center-ci-{int(time.time())}"
    health_url = f"http://127.0.0.1:{args.port}/_stcore/health"

    try:
        run(["docker", "build", "--pull", "-t", args.image, "."])
        remove_container(container_name)
        run(
            [
                "docker",
                "run",
                "-d",
                "--name",
                container_name,
                "-p",
                f"{args.port}:8501",
                args.image,
            ]
        )
        return wait_for_health(container_name, health_url, args.attempts, args.delay)
    finally:
        remove_container(container_name)


if __name__ == "__main__":
    raise SystemExit(main())