test.py

import os
import sys
import argparse
import logging
from pathlib import Path

import requests


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger("colorize_test")


def normalize_space_url(url: str) -> str:
    """
    Accept either a runtime URL (https://<owner>-<space>.hf.space)
    or a repo URL (https://huggingface.co/spaces/<Owner>/<Space>?logs=container)
    and return a runtime base URL suitable for API calls.
    """
    if not url:
        return url
    url = url.strip()
    # Strip query params
    if "?" in url:
        url = url.split("?", 1)[0]
    # Already runtime URL
    if ".hf.space" in url:
        return url.rstrip("/")
    # Convert repo URL to runtime URL
    marker = "/spaces/"
    if "huggingface.co" in url and marker in url:
        try:
            _, rest = url.split(marker, 1)
            owner, space = rest.split("/", 1)
            runtime = f"https://{owner.lower()}-{space.lower()}.hf.space"
            return runtime.rstrip("/")
        except Exception:
            return url.rstrip("/")
    return url.rstrip("/")


def main() -> None:
    parser = argparse.ArgumentParser(description="Test Colorize API")
    parser.add_argument(
        "--image",
        "-i",
        type=str,
        required=False,
        default=r"C:\projects\colorize_text\pexels-andrey-grushnikov-223358-707676.jpg",
        help="Path to input image file",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=180,
        help="Request timeout in seconds (default: 180)",
    )
    args = parser.parse_args()

    base_url = os.getenv(
        "BASE_URL",
        os.getenv(
            "SPACE_HOST",
            "http://localhost:7860",
        ),
    )
    base_url = normalize_space_url(base_url)
    app_check_token = os.getenv("APP_CHECK_TOKEN", "").strip()

    image_path = Path(args.image)
    if not image_path.exists():
        logger.error("Image file not found at: %s", str(image_path))
        sys.exit(1)

    # Health check
    try:
        logger.info("Using BASE_URL: %s", base_url)
        health = requests.get(f"{base_url}/health", timeout=15)
        logger.info("Health check %s: %s", health.status_code, health.text)
    except Exception as e:
        logger.warning("Health check failed: %s", str(e))

    headers = {}
    if app_check_token:
        headers["X-Firebase-AppCheck"] = app_check_token

    files = {"file": (image_path.name, open(image_path, "rb"), "application/octet-stream")}

    logger.info("Sending request to %s", f"{base_url}/colorize")
    try:
        resp = requests.post(
            f"{base_url}/colorize",
            headers=headers,
            files=files,
            timeout=args.timeout,
        )
    finally:
        files["file"][1].close()

    logger.info("Response status: %s", resp.status_code)
    if not resp.ok:
        logger.error("Error response: %s", resp.text)
        sys.exit(2)

    data = resp.json()
    download_url = data.get("download_url")
    api_download_url = data.get("api_download_url")
    filename = data.get("filename")

    logger.info("Colorization success: %s", data.get("success"))
    logger.info("Download URL: %s", download_url)
    logger.info("API Download URL: %s", api_download_url)
    logger.info("Filename: %s", filename)

    # Optionally download the result locally
    if download_url:
        try:
            out_path = Path("colorized_result.jpg")
            r = requests.get(download_url, timeout=60)
            if r.ok:
                out_path.write_bytes(r.content)
                logger.info("Saved colorized image to: %s", str(out_path.resolve()))
            else:
                logger.warning("Could not download from public URL, status: %s", r.status_code)
        except Exception as e:
            logger.warning("Download via public URL failed: %s", str(e))


if __name__ == "__main__":
    main()