test_api.py

"""
test_api.py
───────────
Local ya deployed HF Space ko test karne ke liye.

Usage:
    # Local test (app.py chal raha ho)
    python test_api.py --url http://localhost:7860 --secret your_secret --image logo_test.jpg

    # HF Space test
    python test_api.py --url https://your-space.hf.space --secret your_secret --image logo_test.jpg
"""

import argparse
import base64
import json
import os
import sys
import time

import requests


def test_health(base_url: str) -> bool:
    print(f"[1/3] Health check → {base_url}/health")
    try:
        r = requests.get(f"{base_url}/health", timeout=10)
        data = r.json()
        print(f"      Status: {data.get('status')} | Model loaded: {data.get('model_loaded')}")
        return data.get("status") == "ok"
    except Exception as e:
        print(f"      FAILED: {e}")
        return False


def test_auth_reject(base_url: str) -> None:
    print("[2/3] Auth rejection test (wrong secret)...")
    try:
        r = requests.post(
            f"{base_url}/extract-logo",
            headers={"X-Custom-Secret": "wrong_secret_12345"},
            files={"image": ("test.jpg", b"\xff\xd8\xff", "image/jpeg")},
            timeout=10,
        )
        if r.status_code == 401:
            print("      ✓ Correctly rejected (401 Unauthorized)")
        else:
            print(f"      WARNING: Expected 401, got {r.status_code}")
    except Exception as e:
        print(f"      FAILED: {e}")


def test_extract(base_url: str, secret: str, image_path: str, fmt: str = "PNG") -> None:
    print(f"[3/3] Logo extraction → {image_path}")

    if not os.path.exists(image_path):
        print(f"      ERROR: Image not found: {image_path}")
        return

    with open(image_path, "rb") as f:
        image_bytes = f.read()

    mime = "image/png" if image_path.lower().endswith(".png") else "image/jpeg"

    t0 = time.perf_counter()
    try:
        r = requests.post(
            f"{base_url}/extract-logo",
            headers={"X-Custom-Secret": secret},
            files={"image": (os.path.basename(image_path), image_bytes, mime)},
            params={"output_format": fmt},
            timeout=30,
        )
    except Exception as e:
        print(f"      FAILED: {e}")
        return

    elapsed = round((time.perf_counter() - t0) * 1000, 1)

    if r.status_code != 200:
        print(f"      ERROR {r.status_code}: {r.text}")
        return

    data = r.json()
    print(f"      ✓ Success in {elapsed} ms (server: {data.get('elapsed_ms')} ms)")
    print(f"      Logos found: {data.get('count', 0)}")

    # Save cropped logos
    for i, logo in enumerate(data.get("logos", [])):
        out_path = f"logo_crop_{i+1}.png"
        img_bytes = base64.b64decode(logo["base64"])
        with open(out_path, "wb") as f:
            f.write(img_bytes)
        print(
            f"      Logo {i+1}: {logo['width']}x{logo['height']}px "
            f"| conf={logo['confidence']} | saved → {out_path}"
        )

    if data.get("count", 0) == 0:
        print("      No logos detected. Try a different image or lower CONF_THRESHOLD.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--url",    default="http://localhost:7860", help="API base URL")
    parser.add_argument("--secret", required=True,                   help="APP_SECRET_KEY value")
    parser.add_argument("--image",  required=True,                   help="Path to test image")
    parser.add_argument("--format", default="PNG",
                        choices=["PNG", "JPEG", "TRANSPARENT"],      help="Output format")
    args = parser.parse_args()

    base_url = args.url.rstrip("/")
    ok = test_health(base_url)
    test_auth_reject(base_url)
    if ok:
        test_extract(base_url, args.secret, args.image, args.format)
    else:
        print("\nSkipping extraction test — server not healthy.")
        sys.exit(1)

    print("\nDone!")