ai.allucanget.biz/backend/tests/test_images.py

"""Tests for image upload and retrieval endpoints."""
import io
import os
import pytest
import pytest_asyncio
from httpx import AsyncClient, ASGITransport

from app.main import app
from app import db as db_module

os.environ.setdefault("JWT_SECRET", "test-secret-key-for-testing-only")
# Use a temp dir so file I/O works without polluting project data/
os.environ.setdefault("UPLOAD_DIR", "/tmp/test_uploads")


@pytest.fixture(autouse=True)
def fresh_db():
    db_module._conn = None
    db_module.init_db(":memory:")
    yield
    db_module.close_db()
    db_module._conn = None


@pytest_asyncio.fixture
async def client(fresh_db):
    transport = ASGITransport(app=app)
    async with AsyncClient(transport=transport, base_url="http://test") as ac:
        yield ac


async def _user_token(client) -> str:
    await client.post("/auth/register", json={"email": "user@example.com", "password": "secret123"})
    resp = await client.post("/auth/login", json={"email": "user@example.com", "password": "secret123"})
    return resp.json()["access_token"]


async def _other_token(client) -> str:
    await client.post("/auth/register", json={"email": "other@example.com", "password": "secret123"})
    resp = await client.post("/auth/login", json={"email": "other@example.com", "password": "secret123"})
    return resp.json()["access_token"]


def _png_bytes() -> bytes:
    """Minimal valid 1x1 PNG."""
    return (
        b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01"
        b"\x08\x02\x00\x00\x00\x90wS\xde\x00\x00\x00\x0cIDATx\x9cc\xf8\x0f\x00"
        b"\x00\x01\x01\x00\x05\x18\xd4n\x00\x00\x00\x00IEND\xaeB`\x82"
    )


# ---------------------------------------------------------------------------
# POST /images/upload
# ---------------------------------------------------------------------------

async def test_upload_image_success(client):
    token = await _user_token(client)
    resp = await client.post(
        "/images/upload",
        files={"file": ("test.png", io.BytesIO(_png_bytes()), "image/png")},
        headers={"Authorization": f"Bearer {token}"},
    )
    assert resp.status_code == 201
    data = resp.json()
    assert data["filename"] == "test.png"
    assert data["content_type"] == "image/png"
    assert "id" in data
    assert data["size_bytes"] > 0


async def test_upload_image_unauthenticated(client):
    resp = await client.post(
        "/images/upload",
        files={"file": ("test.png", io.BytesIO(_png_bytes()), "image/png")},
    )
    assert resp.status_code == 401


async def test_upload_image_unsupported_type(client):
    token = await _user_token(client)
    resp = await client.post(
        "/images/upload",
        files={"file": ("doc.pdf", io.BytesIO(
            b"%PDF-fake"), "application/pdf")},
        headers={"Authorization": f"Bearer {token}"},
    )
    assert resp.status_code == 415


async def test_upload_image_too_large(client, monkeypatch):
    import app.routers.images as images_mod
    monkeypatch.setattr(images_mod, "MAX_SIZE_BYTES", 5)
    token = await _user_token(client)
    resp = await client.post(
        "/images/upload",
        files={"file": ("big.png", io.BytesIO(b"x" * 10), "image/png")},
        headers={"Authorization": f"Bearer {token}"},
    )
    assert resp.status_code == 413


# ---------------------------------------------------------------------------
# GET /images/
# ---------------------------------------------------------------------------

async def test_list_images_empty(client):
    token = await _user_token(client)
    resp = await client.get("/images/", headers={"Authorization": f"Bearer {token}"})
    assert resp.status_code == 200
    assert resp.json() == []


async def test_list_images_returns_own_only(client):
    token = await _user_token(client)
    other = await _other_token(client)

    # Upload one image as user, one as other
    for tok, name in [(token, "mine.png"), (other, "theirs.png")]:
        await client.post(
            "/images/upload",
            files={"file": (name, io.BytesIO(_png_bytes()), "image/png")},
            headers={"Authorization": f"Bearer {tok}"},
        )

    resp = await client.get("/images/", headers={"Authorization": f"Bearer {token}"})
    assert resp.status_code == 200
    items = resp.json()
    assert len(items) == 1
    assert items[0]["filename"] == "mine.png"


async def test_list_images_unauthenticated(client):
    resp = await client.get("/images/")
    assert resp.status_code == 401


# ---------------------------------------------------------------------------
# GET /images/{id}/file
# ---------------------------------------------------------------------------

async def test_serve_image_success(client):
    token = await _user_token(client)
    up = await client.post(
        "/images/upload",
        files={"file": ("pixel.png", io.BytesIO(_png_bytes()), "image/png")},
        headers={"Authorization": f"Bearer {token}"},
    )
    image_id = up.json()["id"]

    resp = await client.get(
        f"/images/{image_id}/file",
        headers={"Authorization": f"Bearer {token}"},
    )
    assert resp.status_code == 200
    assert resp.headers["content-type"].startswith("image/png")
    assert resp.content == _png_bytes()


async def test_serve_image_wrong_user(client):
    token = await _user_token(client)
    other = await _other_token(client)

    up = await client.post(
        "/images/upload",
        files={"file": ("secret.png", io.BytesIO(_png_bytes()), "image/png")},
        headers={"Authorization": f"Bearer {token}"},
    )
    image_id = up.json()["id"]

    resp = await client.get(
        f"/images/{image_id}/file",
        headers={"Authorization": f"Bearer {other}"},
    )
    assert resp.status_code == 403


async def test_serve_image_not_found(client):
    token = await _user_token(client)
    resp = await client.get(
        "/images/00000000-0000-0000-0000-000000000000/file",
        headers={"Authorization": f"Bearer {token}"},
    )
    assert resp.status_code == 404