webmail-client/backend/tests/indexer/test_attachments_once.py

"""Attachments are kept once — inside their message (SPEC App. D §D.6 item 1,
acceptance §D.9.5): nothing is stored apart, every way of reading an
attachment returns the exact bytes, and the clean-up pass removes the copies
stored before only when they can be taken out of their message again."""

from __future__ import annotations

import io
import os
import shutil
import time
import uuid
import zipfile
from collections.abc import AsyncIterator
from dataclasses import dataclass
from datetime import UTC, datetime
from email.message import EmailMessage
from pathlib import Path
from types import SimpleNamespace

import httpx
import pytest_asyncio
from httpx import ASGITransport
from PIL import Image
from sqlalchemy import delete, select, update
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.pool import NullPool

from app.config import get_settings
from app.identity.passwords import hash_password
from app.identity.sessions import create_session
from app.indexer import attachments as att_store
from app.indexer import blobstore
from app.indexer.indexer import persist_message
from app.main import create_app
from app.models import Account, Attachment, ComposeAttachment, Mailbox, Message, User
from app.providers.base import RemoteMessage


def _png(w: int, h: int) -> bytes:
    out = io.BytesIO()
    Image.new("RGB", (w, h), (200, 80, 40)).save(out, "PNG")
    return out.getvalue()


PNG = _png(640, 360)
VIDEO = os.urandom(300_000)  # stands in for an .mp4: only the bytes matter here
NOTES = "Zażółć gęślą jaźń\n".encode() * 50
LOGO = _png(32, 32)


def _raw() -> bytes:
    m = EmailMessage()
    m["Subject"] = "Files"
    m["From"] = "Alice <alice@example.com>"
    m["To"] = "me@example.com"
    m["Message-ID"] = f"<{uuid.uuid4().hex}@x.com>"
    m["Date"] = "Mon, 09 Mar 2026 08:36:37 +0100"
    m.set_content("see attached")
    m.add_alternative('<p>see <img src="cid:logo@x"> attached</p>', subtype="html")
    m.get_payload()[1].add_related(LOGO, "image", "png", cid="<logo@x>")
    m.add_attachment(PNG, "image", "png", filename="photo.png")
    m.add_attachment(VIDEO, "video", "mp4", filename="clip.mp4")
    m.add_attachment(NOTES, "text", "plain", filename="notatki.txt")
    return m.as_bytes()


@dataclass
class Ctx:
    client: httpx.AsyncClient
    db: AsyncSession
    user_id: uuid.UUID
    message: Message
    by_name: dict[str, Attachment]


def _user_dir(user_id: uuid.UUID) -> Path:
    return Path(get_settings().blob_root) / str(user_id)


def _stored_bytes(user_id: uuid.UUID) -> int:
    """Bytes in the user's blob store (not the thumbnail or take-out caches)."""
    return sum(
        p.stat().st_size
        for p in _user_dir(user_id).rglob("*")
        if p.is_file() and p.parent.name not in ("thumbs", "cache")
    )


@pytest_asyncio.fixture
async def ctx() -> AsyncIterator[Ctx]:
    engine = create_async_engine(get_settings().database_url, poolclass=NullPool)
    db = async_sessionmaker(engine, expire_on_commit=False)()
    user = User(handle=f"att{uuid.uuid4().hex[:8]}@test.local", password_hash=hash_password("pw"))
    db.add(user)
    await db.flush()
    account = Account(
        user_id=user.id,
        email="me@example.com",
        imap_config={},
        smtp_config={},
        secret_ct=b"",
        discovery_src="manual",
    )
    db.add(account)
    await db.flush()
    inbox = Mailbox(account_id=account.id, user_id=user.id, path="INBOX", special_use="inbox")
    db.add(inbox)
    await db.commit()
    remote = RemoteMessage(uid=1, raw=_raw(), flags=frozenset(), internal_date=datetime.now(UTC))
    result = await persist_message(
        db,
        user_id=user.id,
        account_id=account.id,
        mailbox_id=inbox.id,
        mailbox_special_use="inbox",
        remote=remote,
    )
    await db.commit()
    assert result is not None
    rows = (
        await db.scalars(select(Attachment).where(Attachment.message_id == result.message.id))
    ).all()
    token = await create_session(user.id, "pytest")
    client = httpx.AsyncClient(
        transport=ASGITransport(app=create_app()),
        base_url="http://testserver",
        cookies={"session": token, "csrf": "t"},
    )
    try:
        yield Ctx(client, db, user.id, result.message, {a.filename: a for a in rows})
    finally:
        await client.aclose()
        await db.execute(delete(User).where(User.id == user.id))
        await db.commit()
        await db.close()
        await engine.dispose()
        shutil.rmtree(_user_dir(user.id), ignore_errors=True)


def _url(c: Ctx, name: str, suffix: str = "") -> str:
    return f"/api/messages/{c.message.id}/attachments/{c.by_name[name].id}{suffix}"


async def test_only_the_message_is_stored(ctx: Ctx) -> None:
    assert set(ctx.by_name) == {"photo.png", "clip.mp4", "notatki.txt", "attachment"}
    for att in ctx.by_name.values():
        assert att.part_index is not None and att.part_index >= 0
        assert not blobstore.exists(ctx.user_id, att.blob_sha256)  # no copy of its own
    assert blobstore.exists(ctx.user_id, ctx.message.blob_sha256)
    # §D.9.5: at most 60% of what keeping the copies too took (base64 inside
    # the message is ~1.37x the file; kept apart as well it was ~2.37x).
    raw = _stored_bytes(ctx.user_id)
    before = raw + len(PNG) + len(VIDEO) + len(NOTES) + len(LOGO)
    assert raw <= 0.60 * before


async def test_every_way_of_reading_gives_the_exact_bytes(ctx: Ctx) -> None:
    for name, data in (("photo.png", PNG), ("clip.mp4", VIDEO), ("notatki.txt", NOTES)):
        resp = await ctx.client.get(_url(ctx, name), params={"download": "true"})
        assert resp.status_code == 200 and resp.content == data, name
    ranged = await ctx.client.get(_url(ctx, "clip.mp4"), headers={"Range": "bytes=1000-1999"})
    assert ranged.status_code == 206 and ranged.content == VIDEO[1000:2000]
    inline = await ctx.client.get(f"/api/messages/{ctx.message.id}/cid/logo@x")
    assert inline.status_code == 200 and inline.content == LOGO
    zipped = await ctx.client.get(f"/api/messages/{ctx.message.id}/attachments.zip")
    archive = zipfile.ZipFile(io.BytesIO(zipped.content))
    assert archive.read("clip.mp4") == VIDEO and archive.read("notatki.txt") == NOTES
    thumb = await ctx.client.get(_url(ctx, "photo.png", "/thumb"))
    assert thumb.status_code == 200 and thumb.headers["content-type"] == "image/webp"
    # The streamed copy is a short-lived cache, not a second store.
    assert (_user_dir(ctx.user_id) / "cache" / ctx.by_name["clip.mp4"].blob_sha256).exists()
    assert not blobstore.exists(ctx.user_id, ctx.by_name["clip.mp4"].blob_sha256)


async def test_forwarding_gives_the_files_a_blob_of_their_own(ctx: Ctx) -> None:
    from app.outbox.attachments import load_attachments

    resp = await ctx.client.post(
        f"/api/messages/{ctx.message.id}/forward-attachments", headers={"X-CSRF": "t"}
    )
    assert resp.status_code == 200
    ids = [a["id"] for a in resp.json()]
    assert sorted(a["filename"] for a in resp.json()) == ["clip.mp4", "notatki.txt", "photo.png"]
    draft = SimpleNamespace(payload={"attachment_ids": ids})  # what a send reads
    files = await load_attachments(ctx.db, ctx.user_id, draft)  # type: ignore[arg-type]
    assert {name: data for name, _, data in files} == {
        "photo.png": PNG,
        "clip.mp4": VIDEO,
        "notatki.txt": NOTES,
    }


async def test_old_copies_are_removed_only_when_safe(ctx: Ctx) -> None:
    photo, notes, video = (ctx.by_name[n] for n in ("photo.png", "notatki.txt", "clip.mp4"))
    # As stored before: copies of their own, and no part index.
    old = time.time() - 2 * 3600
    for att, data in ((photo, PNG), (notes, NOTES), (video, VIDEO)):
        blobstore.store(ctx.user_id, data)
        os.utime(blobstore.blob_path(ctx.user_id, att.blob_sha256), (old, old))
    await ctx.db.execute(
        update(Attachment).where(Attachment.message_id == ctx.message.id).values(part_index=None)
    )
    # A file that is not in its message (only its own copy has it) …
    stray = blobstore.store(ctx.user_id, b"only stored apart")
    os.utime(blobstore.blob_path(ctx.user_id, stray), (old, old))
    ctx.db.add(
        Attachment(
            message_id=ctx.message.id,
            user_id=ctx.user_id,
            filename="stray.bin",
            mime_type="application/octet-stream",
            size_bytes=17,
            blob_sha256=stray,
        )
    )
    # … and one a draft is about to send (a compose attachment shares the copy).
    ctx.db.add(
        ComposeAttachment(
            user_id=ctx.user_id,
            filename="notatki.txt",
            mime_type="text/plain",
            size_bytes=len(NOTES),
            blob_sha256=notes.blob_sha256,
        )
    )
    await ctx.db.commit()

    stats = await att_store.dedupe_existing(user_id=ctx.user_id)
    assert stats["found"] == 4 and stats["not_found"] == 1
    assert stats["files_removed"] == 2  # the photo's and the video's copies
    assert not blobstore.exists(ctx.user_id, photo.blob_sha256)
    assert not blobstore.exists(ctx.user_id, video.blob_sha256)
    assert blobstore.exists(ctx.user_id, notes.blob_sha256)  # the draft still needs it
    assert blobstore.exists(ctx.user_id, stray)  # its only copy
    indexes = dict(
        (
            await ctx.db.execute(
                select(Attachment.filename, Attachment.part_index).where(
                    Attachment.message_id == ctx.message.id
                )
            )
        ).all()
    )
    assert indexes["stray.bin"] == att_store.NOT_IN_MESSAGE
    assert all(i is not None and i >= 0 for n, i in indexes.items() if n != "stray.bin")
    # Everything still reads back exactly; a second pass has nothing to do.
    for name, data in (("photo.png", PNG), ("clip.mp4", VIDEO), ("notatki.txt", NOTES)):
        assert (await ctx.client.get(_url(ctx, name), params={"download": "true"})).content == data
    assert (await att_store.dedupe_existing(user_id=ctx.user_id))["rows"] == 0