test_api.py

"""
test_api.py  —  Verify your deployed Space handles both images and PDFs.

Usage:
    python test_api.py --url https://<your-space>.hf.space --image note.jpg
    python test_api.py --url https://<your-space>.hf.space --image note.jpg --pdf doc.pdf
"""

import argparse
import base64
import json
import requests


def pretty(data: dict):
    print(json.dumps(data, indent=2))


def test_health(base_url: str):
    print("\n── Health check ──────────────────────────────────")
    r = requests.get(f"{base_url}/health", timeout=30)
    print(f"Status: {r.status_code}")
    pretty(r.json())
    assert r.json()["status"] == "ok"
    print("✓ PASS")


def test_image_file(base_url: str, path: str):
    print("\n── /ocr/file  [image upload] ─────────────────────")
    with open(path, "rb") as f:
        r = requests.post(f"{base_url}/ocr/file", files={"file": f}, timeout=120)
    print(f"Status: {r.status_code}")
    data = r.json()
    pretty(data)
    assert data.get("success") and data.get("file_type") == "image"
    print(f"✓ PASS  →  '{data['text']}'")


def test_image_base64(base_url: str, path: str):
    print("\n── /ocr/base64  [image base64] ───────────────────")
    with open(path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode()
    r = requests.post(f"{base_url}/ocr/base64",
                      json={"file": b64, "filename": path}, timeout=120)
    print(f"Status: {r.status_code}")
    data = r.json()
    pretty(data)
    assert data.get("success") and data.get("file_type") == "image"
    print(f"✓ PASS  →  '{data['text']}'")


def test_pdf_file(base_url: str, path: str):
    print("\n── /ocr/file  [PDF upload] ───────────────────────")
    with open(path, "rb") as f:
        r = requests.post(f"{base_url}/ocr/file", files={"file": f}, timeout=300)
    print(f"Status: {r.status_code}")
    data = r.json()
    pretty(data)
    assert data.get("success") and data.get("file_type") == "pdf"
    print(f"✓ PASS  →  {data['page_count']} page(s)")


def test_pdf_base64(base_url: str, path: str):
    print("\n── /ocr/base64  [PDF base64] ─────────────────────")
    with open(path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode()
    r = requests.post(f"{base_url}/ocr/base64",
                      json={"file": b64, "filename": path}, timeout=300)
    print(f"Status: {r.status_code}")
    data = r.json()
    pretty(data)
    assert data.get("success") and data.get("file_type") == "pdf"
    print(f"✓ PASS  →  {data['page_count']} page(s)")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--url",   default="http://localhost:7860")
    parser.add_argument("--image", required=True, help="Path to a handwritten image")
    parser.add_argument("--pdf",   default=None,  help="Path to a handwritten PDF (optional)")
    args = parser.parse_args()

    base = args.url.rstrip("/")

    test_health(base)
    test_image_file(base, args.image)
    test_image_base64(base, args.image)

    if args.pdf:
        test_pdf_file(base, args.pdf)
        test_pdf_base64(base, args.pdf)
    else:
        print("\n(Skipping PDF tests — pass --pdf <file> to test PDF support)")

    print("\n\nAll tests passed ✓")