"""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")