test_app.py

"""Tests for the HF Space proxy image endpoints (POST /image, GET /image/{path})."""

import base64
import io
import os
import struct
import tempfile

import pytest
from unittest.mock import patch, MagicMock

# Set environment before importing app
os.environ.setdefault("HF_TOKEN", "test-token")
os.environ.setdefault("HF_PROXY_SECRET", "test-secret")

from fastapi.testclient import TestClient
from PIL import Image

from app import app

client = TestClient(app)
AUTH = {"Authorization": "Bearer test-secret"}


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_png(w: int = 100, h: int = 100) -> bytes:
    """Create a minimal PNG image of the given dimensions."""
    img = Image.new("RGB", (w, h), color=(255, 0, 0))
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue()


def _make_jpeg(w: int = 100, h: int = 100) -> bytes:
    """Create a minimal JPEG image of the given dimensions."""
    img = Image.new("RGB", (w, h), color=(0, 255, 0))
    buf = io.BytesIO()
    img.save(buf, format="JPEG")
    return buf.getvalue()


def _make_jpeg_with_exif(w: int = 100, h: int = 100) -> bytes:
    """Create a JPEG with EXIF metadata (ImageDescription tag)."""
    img = Image.new("RGB", (w, h), color=(0, 0, 255))
    buf = io.BytesIO()
    # Build a minimal EXIF segment with an ImageDescription tag
    # Tag 0x010E (270) = ImageDescription
    from PIL.ExifTags import Base as ExifBase
    from PIL import ExifData  # type: ignore[attr-defined]

    exif_data = img.getexif()
    exif_data[270] = "Test EXIF description"
    img.save(buf, format="JPEG", exif=exif_data.tobytes())
    return buf.getvalue()


# ---------------------------------------------------------------------------
# POST /image tests
# ---------------------------------------------------------------------------


@patch("app.batch_bucket_files")
def test_upload_image_png_converts_to_webp(mock_batch):
    """PNG upload should convert to WebP and return correct metadata."""
    png_bytes = _make_png(200, 200)
    resp = client.post(
        "/image",
        json={
            "session_id": "sess-1",
            "image_data": base64.b64encode(png_bytes).decode(),
            "media_type": "image/png",
        },
        headers=AUTH,
    )
    assert resp.status_code == 200
    data = resp.json()
    assert data["ok"] is True
    assert data["serve_url"].startswith("/image/images/sess-1/")
    assert data["serve_url"].endswith(".webp")
    assert data["width"] <= 200 and data["height"] <= 200
    assert data["size_bytes"] > 0
    assert data["original_size_bytes"] == len(png_bytes)
    mock_batch.assert_called_once()


@patch("app.batch_bucket_files")
def test_upload_image_jpeg_converts_to_webp(mock_batch):
    """JPEG upload should convert to WebP and return success."""
    jpeg_bytes = _make_jpeg(150, 150)
    resp = client.post(
        "/image",
        json={
            "session_id": "sess-2",
            "image_data": base64.b64encode(jpeg_bytes).decode(),
            "media_type": "image/jpeg",
        },
        headers=AUTH,
    )
    assert resp.status_code == 200
    data = resp.json()
    assert data["ok"] is True
    assert data["serve_url"].endswith(".webp")


def test_upload_image_rejects_invalid_magic_bytes():
    """Non-image data should be rejected with 400."""
    resp = client.post(
        "/image",
        json={
            "session_id": "sess-3",
            "image_data": base64.b64encode(b"not-an-image-data-here").decode(),
            "media_type": "application/octet-stream",
        },
        headers=AUTH,
    )
    assert resp.status_code == 400
    assert "Invalid image format" in resp.text


def test_upload_image_rejects_too_large():
    """Images over 20MB should be rejected with 400."""
    # Create data with valid PNG magic bytes but oversized
    oversized = b"\x89PNG" + b"\x00" * 21_000_000
    resp = client.post(
        "/image",
        json={
            "session_id": "sess-4",
            "image_data": base64.b64encode(oversized).decode(),
            "media_type": "image/png",
        },
        headers=AUTH,
    )
    assert resp.status_code == 400
    assert "too large" in resp.text.lower()


@patch("app.batch_bucket_files")
def test_upload_image_resizes_large_image(mock_batch):
    """Images larger than 1920px should be resized with aspect ratio preserved."""
    png_bytes = _make_png(3000, 2000)
    resp = client.post(
        "/image",
        json={
            "session_id": "sess-5",
            "image_data": base64.b64encode(png_bytes).decode(),
            "media_type": "image/png",
        },
        headers=AUTH,
    )
    assert resp.status_code == 200
    data = resp.json()
    assert data["width"] <= 1920
    assert data["height"] <= 1920
    # Check aspect ratio preserved (3000:2000 = 1.5)
    ratio = data["width"] / data["height"]
    assert abs(ratio - 1.5) < 0.02


@patch("app.batch_bucket_files")
def test_upload_image_strips_exif(mock_batch):
    """EXIF metadata should be stripped from the output WebP."""
    try:
        jpeg_bytes = _make_jpeg_with_exif(200, 200)
    except Exception:
        # Fallback: create a simple JPEG if EXIF injection fails
        jpeg_bytes = _make_jpeg(200, 200)

    resp = client.post(
        "/image",
        json={
            "session_id": "sess-6",
            "image_data": base64.b64encode(jpeg_bytes).decode(),
            "media_type": "image/jpeg",
        },
        headers=AUTH,
    )
    assert resp.status_code == 200

    # Capture the bytes that were uploaded to the bucket
    call_args = mock_batch.call_args
    uploaded_bytes = call_args[1]["add"][0][0]

    # Open the uploaded WebP and check for EXIF
    result_img = Image.open(io.BytesIO(uploaded_bytes))
    exif_data = result_img.info.get("exif", b"")
    assert len(exif_data) == 0, "EXIF data should be stripped from output"


def test_upload_image_requires_auth():
    """POST /image without auth should be rejected."""
    png_bytes = _make_png(50, 50)
    resp = client.post(
        "/image",
        json={
            "session_id": "sess-7",
            "image_data": base64.b64encode(png_bytes).decode(),
            "media_type": "image/png",
        },
    )
    assert resp.status_code in (401, 422)


# ---------------------------------------------------------------------------
# GET /image/{path} tests
# ---------------------------------------------------------------------------


def _mock_download_webp(bucket_id, files, token):
    """Mock download that writes a small WebP file to the destination path."""
    dest_path = files[0][1]
    img = Image.new("RGB", (10, 10), color=(128, 128, 128))
    img.save(dest_path, format="WEBP")


@patch("app.download_bucket_files", side_effect=_mock_download_webp)
def test_serve_image_returns_webp_bytes(mock_dl):
    """GET /image should return WebP bytes with correct headers."""
    resp = client.get("/image/images/sess123/abc.webp")
    assert resp.status_code == 200
    assert resp.headers["content-type"] == "image/webp"
    assert "max-age=86400" in resp.headers.get("cache-control", "")


@patch("app.download_bucket_files", side_effect=_mock_download_webp)
def test_serve_image_no_auth_required(mock_dl):
    """GET /image should work without any Authorization header."""
    resp = client.get("/image/images/sess456/def.webp")
    assert resp.status_code == 200


@patch("app.download_bucket_files", side_effect=Exception("not found"))
def test_serve_image_returns_404_on_missing(mock_dl):
    """GET /image for a nonexistent file should return 404."""
    resp = client.get("/image/images/nonexistent.webp")
    assert resp.status_code == 404


# ---------------------------------------------------------------------------
# Regression: existing endpoints still work
# ---------------------------------------------------------------------------


def test_health_endpoint_still_works():
    """GET /health should still return ok."""
    resp = client.get("/health")
    assert resp.status_code == 200
    assert resp.json()["status"] == "ok"


@patch("app.batch_bucket_files")
def test_upload_image_gif_input(mock_batch):
    """GIF input should be accepted and converted to WebP."""
    img = Image.new("RGB", (80, 60), color=(255, 255, 0))
    buf = io.BytesIO()
    img.save(buf, format="GIF")
    gif_bytes = buf.getvalue()

    resp = client.post(
        "/image",
        json={
            "session_id": "sess-gif",
            "image_data": base64.b64encode(gif_bytes).decode(),
            "media_type": "image/gif",
        },
        headers=AUTH,
    )
    assert resp.status_code == 200
    data = resp.json()
    assert data["ok"] is True
    assert data["serve_url"].endswith(".webp")