test_api.py

"""
EL HELAL Studio — API Test Suite
Run with: python test_api.py
"""

import requests
import json
import io
import sys
from pathlib import Path
from PIL import Image

BASE_URL = "http://127.0.0.1:9000"

# ── Helpers ──────────────────────────────────────────────────────────────────


def status():
    r = requests.get(f"{BASE_URL}/status")
    r.raise_for_status()
    return r.json()


def wait_for_ai(timeout=60):
    """Poll /status until ai_ready is True."""
    import time

    for _ in range(timeout):
        s = status()
        if s.get("ai_ready"):
            return True
        time.sleep(2)
    raise TimeoutError("AI model did not become ready within timeout")


def create_test_image(width=400, height=600, color=(100, 150, 200)) -> io.BytesIO:
    """Create a synthetic JPEG for testing upload."""
    img = Image.new("RGB", (width, height), color)
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=85)
    buf.seek(0)
    buf.name = "test_portrait.jpg"
    return buf


def upload_image(file_obj) -> dict:
    r = requests.post(f"{BASE_URL}/upload", files={"file": file_obj})
    r.raise_for_status()
    return r.json()


# ── Test Cases ────────────────────────────────────────────────────────────────


def test_status():
    print("\n[TEST] GET /status")
    data = status()
    assert "ai_ready" in data, f"Expected 'ai_ready' key, got: {data}"
    print(f"  → ai_ready: {data['ai_ready']}")


def test_settings_get():
    print("\n[TEST] GET /settings")
    r = requests.get(f"{BASE_URL}/settings")
    r.raise_for_status()
    data = r.json()
    print(f"  → settings: {json.dumps(data, indent=4)}")


def test_settings_update():
    print("\n[TEST] POST /settings")
    payload = {"retouch": {"enabled": True, "sensitivity": 2.5, "tone_smoothing": 0.5}}
    r = requests.post(f"{BASE_URL}/settings", json=payload)
    r.raise_for_status()
    data = r.json()
    assert data.get("status") == "success"
    print(f"  → status: {data['status']}")


def test_frames_list():
    print("\n[TEST] GET /frames")
    r = requests.get(f"{BASE_URL}/frames")
    r.raise_for_status()
    data = r.json()
    assert "frames" in data
    print(f"  → frames count: {len(data['frames'])}")


def test_upload():
    print("\n[TEST] POST /upload")
    img = create_test_image()
    data = upload_image(img)
    assert "id" in data
    assert "thumb_url" in data
    print(f"  → file_id: {data['id']}")
    print(f"  → thumb_url: {data['thumb_url']}")
    return data  # Return full dict for caller


def test_process_image(file_id: str):
    print(f"\n[TEST] POST /process/{file_id}")
    payload = {
        "name": "Test User",
        "id_number": "12345",
        "do_restore": False,
        "fidelity": 0.5,
        "do_rmbg": True,
        "do_color": True,
        "do_retouch": True,
        "do_crop": True,
        "add_studio_name": True,
        "add_logo": True,
        "add_date": True,
    }
    r = requests.post(f"{BASE_URL}/process/{file_id}", data=payload)
    r.raise_for_status()
    data = r.json()
    assert "result_url" in data
    print(f"  → result_url: {data['result_url']}")
    print(f"  → preview_url: {data['preview_url']}")


def test_process_image_with_restore(file_id: str):
    print(f"\n[TEST] POST /process/{file_id} (with restoration)")
    payload = {
        "name": "Restored User",
        "id_number": "99999",
        "do_restore": True,
        "fidelity": 0.6,
        "do_rmbg": True,
        "do_color": True,
        "do_retouch": False,
        "do_crop": True,
        "add_studio_name": True,
        "add_logo": False,
        "add_date": False,
    }
    r = requests.post(f"{BASE_URL}/process/{file_id}", data=payload)
    r.raise_for_status()
    data = r.json()
    print(f"  → result_url: {data['result_url']}")


def test_process_manual_crop(file_id: str):
    print(f"\n[TEST] POST /process/{file_id} (manual crop)")
    payload = {
        "name": "Manual Crop",
        "id_number": "55555",
        "x1": 100,
        "y1": 50,
        "x2": 300,
        "y2": 450,
        "do_crop": True,
        "do_rmbg": True,
        "do_color": False,
        "do_retouch": False,
        "add_studio_name": False,
        "add_logo": False,
        "add_date": False,
    }
    r = requests.post(f"{BASE_URL}/process/{file_id}", data=payload)
    r.raise_for_status()
    data = r.json()
    print(f"  → result_url: {data['result_url']}")


def test_frames_upload():
    print("\n[TEST] POST /frames (upload a test frame)")
    img = Image.new("RGBA", (200, 200), (200, 100, 100, 255))
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    buf.seek(0)
    r = requests.post(
        f"{BASE_URL}/frames",
        files={"file": ("test_frame.png", buf, "image/png")},
    )
    r.raise_for_status()
    data = r.json()
    assert data.get("status") == "success"
    frame = data["frame"]
    print(f"  → uploaded: {frame['filename']}")
    return frame["filename"]


def test_frames_delete(filename: str):
    print(f"\n[TEST] DELETE /frames/{filename}")
    r = requests.delete(f"{BASE_URL}/frames/{filename}")
    r.raise_for_status()
    data = r.json()
    assert data.get("status") == "success"
    print(f"  → deleted: {filename}")


def test_backup_export():
    print("\n[TEST] POST /backup/export")
    client_data = {"theme": "dark", "last_user": "tester"}
    r = requests.post(f"{BASE_URL}/backup/export", json={"client_data": client_data})
    r.raise_for_status()
    assert r.headers["content-type"].startswith("application/zip")
    print(f"  → Content-Length: {len(r.content)} bytes")
    return r.content


def test_backup_import(zip_bytes: bytes):
    print("\n[TEST] POST /backup/import")
    files = {"file": ("backup.zip", io.BytesIO(zip_bytes), "application/zip")}
    r = requests.post(f"{BASE_URL}/backup/import", files=files)
    r.raise_for_status()
    data = r.json()
    print(f"  → status: {data['status']}")


def test_clear_all():
    print("\n[TEST] POST /clear-all")
    r = requests.post(f"{BASE_URL}/clear-all")
    r.raise_for_status()
    data = r.json()
    assert data.get("status") == "success"
    print(f"  → removed_count: {data.get('removed_count')}")


# ── Main ─────────────────────────────────────────────────────────────────────


def main():
    print("=" * 60)
    print("EL HELAL Studio — API Test Suite")
    print("=" * 60)

    # 1. Status check
    test_status()

    # 2. Wait for AI to be ready
    print("\n[INFO] Waiting for AI models to load...")
    try:
        wait_for_ai(timeout=90)
        print("[INFO] AI models are ready.")
    except TimeoutError as e:
        print(f"[WARN] {e}")

    # 3. Settings
    test_settings_get()
    test_settings_update()

    # 4. Frames
    test_frames_list()

    # 5. Upload + Process pipeline
    upload_data = test_upload()
    file_id = upload_data["id"]

    test_process_image(file_id)
    test_process_image_with_restore(file_id)
    test_process_image_with_restore(file_id)  # extra run for variance

    # 6. Manual crop
    upload2 = test_upload()
    test_process_manual_crop(upload2["id"])

    # 7. Frames upload/delete cycle
    frame_filename = test_frames_upload()
    test_frames_delete(frame_filename)

    # 8. Backup cycle
    zip_bytes = test_backup_export()
    test_backup_import(zip_bytes)

    # 9. Clear all
    test_clear_all()

    print("\n" + "=" * 60)
    print("ALL TESTS PASSED")
    print("=" * 60)


if __name__ == "__main__":
    main()