app/tests/test_photos_api.py

from __future__ import annotations

import base64
import uuid

import pytest
from httpx import AsyncClient
from sqlalchemy.ext.asyncio import AsyncSession

from app.db.models import Document, IngestStatus, Report, ReportStatus


def _tiny_png_bytes() -> bytes:
    # 1x1 PNG
    return base64.b64decode(
        "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO6q3d8AAAAASUVORK5CYII="
    )


@pytest.mark.asyncio
async def test_photo_policy_endpoint_returns_sections(async_client: AsyncClient, test_db: AsyncSession) -> None:
    tenant_id = "tenant_test"
    doc = Document(
        id=str(uuid.uuid4()),
        tenant_id=tenant_id,
        filename="x.pdf",
        file_path="C:/tmp/x.pdf",
        status=IngestStatus.complete,
        survey_level=3,
    )
    rep = Report(
        id=str(uuid.uuid4()),
        tenant_id=tenant_id,
        document_id=doc.id,
        status=ReportStatus.pending,
        survey_level=3,
    )
    test_db.add_all([doc, rep])
    await test_db.commit()

    r = await async_client.get(f"/reports/{rep.id}/photo-policy")
    assert r.status_code == 200
    data = r.json()
    assert data["report_id"] == rep.id
    assert data["survey_level"] == 3
    assert isinstance(data["sections"], list)
    assert any(x["code"] == "D" for x in data["sections"])
    assert "policy_configuration_required" in data
    assert isinstance(data["policy_configuration_required"], bool)
    assert "indexed_upload_count" in data
    assert isinstance(data["indexed_upload_count"], int)
    for row in data["sections"]:
        assert "source" in row
        assert row["source"] in (
            "override",
            "tenant_library",
            "tenant_insufficient",
            "tenant_no_documents",
            "data_driven_disabled",
        )


@pytest.mark.asyncio
async def test_upload_list_and_get_section_photo_roundtrip(async_client: AsyncClient, test_db: AsyncSession, tmp_path) -> None:
    tenant_id = "tenant_test"
    doc = Document(
        id=str(uuid.uuid4()),
        tenant_id=tenant_id,
        filename="x.pdf",
        file_path="C:/tmp/x.pdf",
        status=IngestStatus.complete,
        survey_level=3,
    )
    rep = Report(
        id=str(uuid.uuid4()),
        tenant_id=tenant_id,
        document_id=doc.id,
        status=ReportStatus.pending,
        survey_level=3,
    )
    test_db.add_all([doc, rep])
    await test_db.commit()

    from app.config import settings

    settings.upload_dir = tmp_path / "uploads"
    settings.upload_dir.mkdir(parents=True, exist_ok=True)

    png = _tiny_png_bytes()
    files = [("files", ("p.png", png, "image/png"))]

    up = await async_client.post(f"/reports/{rep.id}/sections/D/photos", files=files)
    assert up.status_code == 200
    saved = up.json()["saved"]
    assert saved and saved[0]["photo_id"]

    lst = await async_client.get(f"/reports/{rep.id}/sections/D/photos")
    assert lst.status_code == 200
    photos = lst.json()["photos"]
    assert len(photos) == 1

    url = photos[0]["url"]
    img = await async_client.get(url)
    assert img.status_code == 200
    assert img.headers["content-type"].startswith("image/")

    from app.api.auth import mint_token

    token, _ = mint_token(tenant_id)
    img_q = await async_client.get(url, params={"access_token": token}, headers={})
    assert img_q.status_code == 200
    assert img_q.headers["content-type"].startswith("image/")

    assert photos[0].get("selected_for_ai") is False

    photo_id = photos[0]["photo_id"]
    deleted = await async_client.delete(f"/reports/{rep.id}/sections/D/photos/{photo_id}")
    assert deleted.status_code == 200
    assert deleted.json()["removed"] is True

    lst2 = await async_client.get(f"/reports/{rep.id}/sections/D/photos")
    assert lst2.status_code == 200
    assert len(lst2.json()["photos"]) == 0


@pytest.mark.asyncio
async def test_photo_ai_selection_max_two(async_client: AsyncClient, test_db: AsyncSession, tmp_path) -> None:
    tenant_id = "tenant_test"
    doc = Document(
        id=str(uuid.uuid4()),
        tenant_id=tenant_id,
        filename="x.pdf",
        file_path="C:/tmp/x.pdf",
        status=IngestStatus.complete,
        survey_level=3,
    )
    rep = Report(
        id=str(uuid.uuid4()),
        tenant_id=tenant_id,
        document_id=doc.id,
        status=ReportStatus.pending,
        survey_level=3,
    )
    test_db.add_all([doc, rep])
    await test_db.commit()

    from app.config import settings

    settings.upload_dir = tmp_path / "uploads"
    settings.upload_dir.mkdir(parents=True, exist_ok=True)
    settings.max_section_photos_per_section = 5
    settings.max_section_photos_for_ai = 2

    png = _tiny_png_bytes()
    ids: list[str] = []
    for i in range(3):
        up = await async_client.post(
            f"/reports/{rep.id}/sections/D/photos",
            files=[("files", (f"p{i}.png", png, "image/png"))],
        )
        assert up.status_code == 200
        ids.append(up.json()["saved"][0]["photo_id"])

    sel = await async_client.put(
        f"/reports/{rep.id}/sections/D/photos/ai-selection",
        json={"photo_ids": ids[:2]},
    )
    assert sel.status_code == 200
    body = sel.json()
    assert body["selected_for_ai_count"] == 2
    selected = {p["photo_id"] for p in body["photos"] if p["selected_for_ai"]}
    assert selected == set(ids[:2])

    too_many = await async_client.put(
        f"/reports/{rep.id}/sections/D/photos/ai-selection",
        json={"photo_ids": ids},
    )
    assert too_many.status_code == 422

    policy = await async_client.get(f"/reports/{rep.id}/photo-policy")
    assert policy.status_code == 200
    limits = policy.json()["photo_limits"]
    assert limits["max_photos_per_section"] == 5
    assert limits["max_photos_for_ai"] == 2


@pytest.mark.asyncio
async def test_upload_rejects_more_than_five_photos(async_client: AsyncClient, test_db: AsyncSession, tmp_path) -> None:
    tenant_id = "tenant_test"
    doc = Document(
        id=str(uuid.uuid4()),
        tenant_id=tenant_id,
        filename="x.pdf",
        file_path="C:/tmp/x.pdf",
        status=IngestStatus.complete,
        survey_level=3,
    )
    rep = Report(
        id=str(uuid.uuid4()),
        tenant_id=tenant_id,
        document_id=doc.id,
        status=ReportStatus.pending,
        survey_level=3,
    )
    test_db.add_all([doc, rep])
    await test_db.commit()

    from app.config import settings

    settings.upload_dir = tmp_path / "uploads"
    settings.upload_dir.mkdir(parents=True, exist_ok=True)
    settings.max_section_photos_per_section = 5

    png = _tiny_png_bytes()
    for _ in range(5):
        r = await async_client.post(
            f"/reports/{rep.id}/sections/D/photos",
            files=[("files", ("p.png", png, "image/png"))],
        )
        assert r.status_code == 200

    sixth = await async_client.post(
        f"/reports/{rep.id}/sections/D/photos",
        files=[("files", ("p6.png", png, "image/png"))],
    )
    assert sixth.status_code == 422