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

"""The nightly clean-up of stored files (app/indexer/gc.py): what nothing uses
goes, everything still used — or possibly about to be — stays."""

from __future__ import annotations

import datetime as dt
import os
import shutil
import time
import uuid
from collections.abc import AsyncIterator
from dataclasses import dataclass
from pathlib import Path

import pytest_asyncio
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.indexer import blobstore, gc
from app.models import Account, Attachment, ComposeAttachment, Draft, Mailbox, Message, User

OLD = time.time() - 2 * 3600


@dataclass
class Ctx:
    db: AsyncSession
    user_id: uuid.UUID
    account_id: uuid.UUID
    mailbox_id: uuid.UUID


@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"gc{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()
    try:
        yield Ctx(db, user.id, account.id, inbox.id)
    finally:
        await db.execute(delete(User).where(User.id == user.id))
        await db.commit()
        await db.close()
        await engine.dispose()
        shutil.rmtree(Path(get_settings().blob_root) / str(user.id), ignore_errors=True)


def _file(c: Ctx, data: bytes, *, old: bool = True) -> str:
    sha = blobstore.store(c.user_id, data)
    if old:
        os.utime(blobstore.blob_path(c.user_id, sha), (OLD, OLD))
    return sha


async def _message(c: Ctx, sha: str, uid: int) -> Message:
    m = Message(
        user_id=c.user_id,
        account_id=c.account_id,
        mailbox_id=c.mailbox_id,
        uid=uid,
        internal_date=dt.datetime.now(dt.UTC),
        blob_sha256=sha,
    )
    c.db.add(m)
    await c.db.commit()
    return m


async def test_unused_files_go_and_used_ones_stay(ctx: Ctx) -> None:
    kept_message = _file(ctx, b"a message still here")
    await _message(ctx, kept_message, 1)
    deleted_message = _file(ctx, b"a message deleted forever")
    fresh = _file(ctx, b"a message being stored right now", old=False)
    legacy = _file(ctx, b"an attachment kept as its own file")
    msg = await _message(ctx, _file(ctx, b"its message"), 2)
    ctx.db.add(
        Attachment(
            message_id=msg.id,
            user_id=ctx.user_id,
            filename="x.bin",
            mime_type="application/octet-stream",
            size_bytes=1,
            blob_sha256=legacy,
            part_index=None,  # not looked for in its message yet
        )
    )
    await ctx.db.commit()
    # Previews: one of an attachment still here, one of a deleted one.
    thumbs = blobstore.blob_path(ctx.user_id, "00").parent.parent / "thumbs"
    thumbs.mkdir(parents=True, exist_ok=True)
    live_thumb, dead_thumb = thumbs / f"{legacy}-480.webp", thumbs / f"{'f' * 64}-480.webp"
    for t in (live_thumb, dead_thumb):
        t.write_bytes(b"webp")
        os.utime(t, (OLD, OLD))
    # A sent message waiting for its Sent copy: only its job refers to it.
    waiting = _file(ctx, b"sent, not yet filed in Sent")
    await gc.hold(ctx.user_id, waiting)

    stats = await gc.sweep(ctx.user_id)

    assert stats["files"] == 1 and stats["thumbs"] == 1
    assert not blobstore.exists(ctx.user_id, deleted_message)
    for sha in (kept_message, fresh, legacy, waiting, msg.blob_sha256):
        assert blobstore.exists(ctx.user_id, sha)
    assert live_thumb.exists() and not dead_thumb.exists()

    # Filed (or given up): no longer held, so it goes on the next night.
    await gc.release(ctx.user_id, waiting)
    assert (await gc.sweep(ctx.user_id))["files"] == 1
    assert not blobstore.exists(ctx.user_id, waiting)


async def test_abandoned_uploads_go_after_a_week(ctx: Ctx) -> None:
    week_ago = dt.datetime.now(dt.UTC) - dt.timedelta(days=8)

    def upload(data: bytes, created: dt.datetime) -> ComposeAttachment:
        row = ComposeAttachment(
            user_id=ctx.user_id,
            filename="f.txt",
            mime_type="text/plain",
            size_bytes=len(data),
            blob_sha256=_file(ctx, data),
        )
        ctx.db.add(row)
        return row

    abandoned = upload(b"closed without saving", week_ago)
    in_draft = upload(b"in a draft still being written", week_ago)
    recent = upload(b"uploaded a minute ago", dt.datetime.now(dt.UTC))
    await ctx.db.flush()
    await ctx.db.execute(
        update(ComposeAttachment)
        .where(ComposeAttachment.id.in_([abandoned.id, in_draft.id]))
        .values(created_at=week_ago)
    )
    ctx.db.add(
        Draft(
            user_id=ctx.user_id,
            account_id=ctx.account_id,
            kind="new",
            payload={"attachment_ids": [str(in_draft.id)]},
        )
    )
    await ctx.db.commit()

    stats = await gc.sweep(ctx.user_id)

    assert stats["uploads"] == 1 and stats["files"] == 1
    left = set(
        await ctx.db.scalars(
            select(ComposeAttachment.id).where(ComposeAttachment.user_id == ctx.user_id)
        )
    )
    assert left == {in_draft.id, recent.id}
    assert not blobstore.exists(ctx.user_id, abandoned.blob_sha256)
    assert blobstore.exists(ctx.user_id, in_draft.blob_sha256)
    assert blobstore.exists(ctx.user_id, recent.blob_sha256)