test_api.py

"""
test_api.py — run against a live server (local or Hugging Face Space).

Usage:
    # against local server
    python test_api.py

    # against HF Space
    python test_api.py --base-url https://<your-space-name>.hf.space
"""

import argparse
import io
import sys
import numpy as np
from PIL import Image

try:
    import requests
except ImportError:
    sys.exit("Install requests first:  pip install requests")

DEFAULT_BASE = "http://localhost:7860"


def make_test_image(width: int = 512, height: int = 300) -> bytes:
    """Generate a synthetic grayscale document-like image."""
    rng = np.random.default_rng(42)
    img = np.full((height, width), 240, dtype=np.uint8)
    # Add some random noise stripes to simulate a degraded document
    for _ in range(30):
        row = rng.integers(0, height)
        img[row, :] = rng.integers(80, 180)
    buf = io.BytesIO()
    Image.fromarray(img, mode="L").save(buf, format="PNG")
    buf.seek(0)
    return buf.read()


def assert_png(data: bytes, label: str) -> None:
    try:
        img = Image.open(io.BytesIO(data))
        assert img.format == "PNG", f"Expected PNG, got {img.format}"
        print(f"  ✓ {label}: {img.size[0]}x{img.size[1]} PNG")
    except Exception as exc:
        print(f"  ✗ {label}: {exc}")
        sys.exit(1)


def test_health(base: str) -> None:
    print("[ 1 ] Health check")
    r = requests.get(f"{base}/health", timeout=30)
    assert r.status_code == 200, f"Expected 200, got {r.status_code}"
    data = r.json()
    assert data.get("status") == "healthy", f"Unexpected body: {data}"
    assert data.get("model_loaded") is True, "Model not loaded!"
    print(f"  ✓ {data}")


def test_enhance(base: str, img_bytes: bytes) -> None:
    print("[ 2 ] POST /enhance")
    r = requests.post(
        f"{base}/enhance",
        files={"file": ("test.png", img_bytes, "image/png")},
        timeout=120,
    )
    assert r.status_code == 200, f"Expected 200, got {r.status_code}: {r.text}"
    assert "image/png" in r.headers.get("content-type", ""), "Response is not PNG"
    assert_png(r.content, "enhanced image")


def test_binarize(base: str, img_bytes: bytes) -> None:
    print("[ 3 ] POST /binarize  (threshold=0.65)")
    r = requests.post(
        f"{base}/binarize",
        params={"threshold": 0.65},
        files={"file": ("test.png", img_bytes, "image/png")},
        timeout=120,
    )
    assert r.status_code == 200, f"Expected 200, got {r.status_code}: {r.text}"
    assert "image/png" in r.headers.get("content-type", ""), "Response is not PNG"
    assert_png(r.content, "binarized image")


def test_invalid_file(base: str) -> None:
    print("[ 4 ] POST /enhance  (invalid file → expect 400)")
    r = requests.post(
        f"{base}/enhance",
        files={"file": ("bad.txt", b"not an image at all", "text/plain")},
        timeout=30,
    )
    assert r.status_code == 400, f"Expected 400, got {r.status_code}: {r.text}"
    print(f"  ✓ correctly rejected with 400: {r.json().get('detail')}")


def test_threshold_validation(base: str, img_bytes: bytes) -> None:
    print("[ 5 ] POST /binarize  (threshold=2.0 → expect 422)")
    r = requests.post(
        f"{base}/binarize",
        params={"threshold": 2.0},
        files={"file": ("test.png", img_bytes, "image/png")},
        timeout=30,
    )
    assert r.status_code == 422, f"Expected 422, got {r.status_code}: {r.text}"
    print(f"  ✓ correctly rejected with 422")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-url", default=DEFAULT_BASE, help="Base URL of the API")
    args = parser.parse_args()
    base = args.base_url.rstrip("/")

    print(f"\nRunning tests against: {base}\n")

    img_bytes = make_test_image()

    test_health(base)
    test_enhance(base, img_bytes)
    test_binarize(base, img_bytes)
    test_invalid_file(base)
    test_threshold_validation(base, img_bytes)

    print("\n✅  All tests passed!")


if __name__ == "__main__":
    main()