simbarag/tests/unit/test_file_upload.py

"""Tests for conversation file attachments: extraction, /upload and stream-query file_key."""

import io
import threading
from datetime import timedelta
from types import SimpleNamespace
from unittest.mock import AsyncMock, patch

import fitz
import pytest
from PIL import Image
from quart import Quart
from quart.datastructures import FileStorage
from quart_jwt_extended import JWTManager, create_access_token

import blueprints.conversation as conv
from utils.file_extract import (
    MAX_ATTACHMENT_CHARS,
    TRUNCATION_MARKER,
    FileExtractionError,
    attachment_block,
    cap_text,
    extract_text,
    file_extension,
)

CONV_ID = "11111111-1111-1111-1111-111111111111"
OTHER_ID = "22222222-2222-2222-2222-222222222222"


def _pdf_bytes(text: str) -> bytes:
    doc = fitz.open()
    page = doc.new_page()
    page.insert_text((72, 72), text)
    data = doc.tobytes()
    doc.close()
    return data


def _png_bytes() -> bytes:
    buf = io.BytesIO()
    Image.new("RGB", (4, 4), "orange").save(buf, format="PNG")
    return buf.getvalue()


class TestFileExtension:
    @pytest.mark.parametrize(
        "name,mime,expected",
        [
            ("a.pdf", "application/pdf", "pdf"),
            ("notes", "text/markdown", "md"),
            ("notes.md", "application/octet-stream", "md"),
            ("data.CSV", "", "csv"),
            ("x.txt", "text/plain; charset=utf-8", "txt"),
            ("x.markdown", None, "md"),
            ("x.docx", "application/vnd.openxmlformats", None),
            ("x.zip", "application/zip", None),
        ],
    )
    def test_detects_supported_types(self, name, mime, expected):
        assert file_extension(name, mime) == expected


class TestExtractText:
    def test_pdf(self):
        assert "Simba eats kibble" in extract_text(
            _pdf_bytes("Simba eats kibble"), "pdf"
        )

    def test_corrupt_pdf_raises(self):
        with pytest.raises(FileExtractionError):
            extract_text(b"not a pdf", "pdf")

    def test_text_decodes_utf8_with_replacement(self):
        assert extract_text("héllo".encode() + b"\xff", "txt") == "héllo�"

    def test_csv_and_md_are_plain_text(self):
        assert extract_text(b"\xef\xbb\xbfa,b\n1,2", "csv") == "a,b\n1,2"
        assert extract_text(b"# Title", "md") == "# Title"


class TestTruncation:
    def test_short_text_untouched(self):
        assert cap_text("abc") == "abc"

    def test_long_text_capped_with_marker(self):
        out = cap_text("x" * (MAX_ATTACHMENT_CHARS + 50))
        assert out.endswith(TRUNCATION_MARKER)
        assert out.count("x") == MAX_ATTACHMENT_CHARS

    def test_attachment_block_wraps_and_caps(self):
        block = attachment_block('a"b.txt', "y" * (MAX_ATTACHMENT_CHARS + 1))
        assert block.startswith('<attached file "a\'b.txt">\n')
        assert block.endswith(f"{TRUNCATION_MARKER}\n</attached>")


class TestFileKeyPrefix:
    def test_own_conversation(self):
        assert conv._file_key_belongs_to(
            f"conversations/{CONV_ID}/files/x.pdf", CONV_ID
        )

    @pytest.mark.parametrize(
        "key",
        [
            f"conversations/{OTHER_ID}/files/x.pdf",
            f"conversations/{CONV_ID}/x.jpg",
            f"conversations/{CONV_ID}/files/../../{OTHER_ID}/files/x.pdf",
            f"conversations/{CONV_ID}",
            None,
        ],
    )
    def test_rejects_other_keys(self, key):
        assert not conv._file_key_belongs_to(key, CONV_ID)


@pytest.fixture
def app():
    app = Quart(__name__)
    app.config["JWT_SECRET_KEY"] = "test-secret"
    app.config["JWT_ACCESS_TOKEN_EXPIRES"] = timedelta(hours=1)
    JWTManager(app)
    app.register_blueprint(conv.conversation_blueprint)
    return app


@pytest.fixture
async def headers(app):
    async with app.app_context():
        token = create_access_token(identity="user-1")
    return {"Authorization": f"Bearer {token}"}


def _owned(conversation_id, user_id):
    if conversation_id == CONV_ID and user_id == "user-1":
        return SimpleNamespace(id=CONV_ID)
    return None


@pytest.fixture
def s3():
    with (
        patch.object(conv, "_get_owned_conversation", AsyncMock(side_effect=_owned)),
        patch.object(conv, "s3_upload_file", AsyncMock()) as upload_file,
        patch.object(conv, "s3_upload_image", AsyncMock()) as upload_image,
    ):
        yield SimpleNamespace(upload_file=upload_file, upload_image=upload_image)


async def _upload(
    app, headers, data: bytes, filename: str, mime: str, conversation_id=CONV_ID
):
    return await app.test_client().post(
        "/api/conversation/upload",
        headers=headers,
        form={"conversation_id": conversation_id},
        files={
            "file": FileStorage(io.BytesIO(data), filename=filename, content_type=mime)
        },
    )


class TestUploadRoute:
    async def test_pdf_is_stored_with_filename(self, app, headers, s3):
        resp = await _upload(
            app, headers, _pdf_bytes("hello"), "Vet bill.pdf", "application/pdf"
        )
        assert resp.status_code == 200
        body = await resp.get_json()
        assert body["kind"] == "file"
        assert body["name"] == "Vet bill.pdf"
        assert body["chars"] >= len("hello")
        assert body["file_key"].startswith(f"conversations/{CONV_ID}/files/")
        assert body["file_key"].endswith(".pdf")
        _, key, mime, metadata = s3.upload_file.await_args.args
        assert key == body["file_key"] and mime == "application/pdf"
        assert metadata == {"filename": "Vet%20bill.pdf"}

    async def test_markdown_by_extension(self, app, headers, s3):
        resp = await _upload(app, headers, b"# hi", "n.md", "application/octet-stream")
        body = await resp.get_json()
        assert resp.status_code == 200
        assert body["chars"] == 4 and body["file_key"].endswith(".md")

    async def test_image_goes_through_process_image(self, app, headers, s3):
        resp = await _upload(app, headers, _png_bytes(), "cat.png", "image/png")
        assert resp.status_code == 200
        body = await resp.get_json()
        assert body["kind"] == "image"
        assert body["image_key"].startswith(f"conversations/{CONV_ID}/")
        s3.upload_image.assert_awaited_once()
        s3.upload_file.assert_not_awaited()

    async def test_image_processing_runs_off_the_event_loop(
        self, app, headers, s3, monkeypatch
    ):
        threads = []
        real = conv.process_image

        def recording(*args):
            threads.append(threading.current_thread())
            return real(*args)

        monkeypatch.setattr(conv, "process_image", recording)
        resp = await _upload(app, headers, _png_bytes(), "cat.png", "image/png")
        assert resp.status_code == 200
        assert threads and threads[0] is not threading.main_thread()

    async def test_unsupported_type_is_415(self, app, headers, s3):
        resp = await _upload(app, headers, b"PK\x03\x04", "a.zip", "application/zip")
        assert resp.status_code == 415
        s3.upload_file.assert_not_awaited()

    async def test_unsupported_image_type_is_415(self, app, headers, s3):
        resp = await _upload(app, headers, b"GIF89a", "a.gif", "image/gif")
        assert resp.status_code == 415

    async def test_corrupt_pdf_is_400(self, app, headers, s3):
        resp = await _upload(app, headers, b"nope", "a.pdf", "application/pdf")
        assert resp.status_code == 400

    async def test_pdf_without_text_layer_is_422(self, app, headers, s3):
        doc = fitz.open()
        doc.new_page()  # blank page: no text, like a scan without OCR
        resp = await _upload(app, headers, doc.tobytes(), "scan.pdf", "application/pdf")
        assert resp.status_code == 422
        assert "no text layer" in (await resp.get_json())["error"]
        s3.upload_file.assert_not_awaited()

    async def test_empty_text_file_is_422(self, app, headers, s3):
        resp = await _upload(app, headers, b"  \n", "empty.txt", "text/plain")
        assert resp.status_code == 422
        s3.upload_file.assert_not_awaited()

    async def test_extraction_runs_off_the_event_loop(self, app, headers, s3):
        calls = []

        async def fake_to_thread(fn, *args):
            calls.append(fn)
            return fn(*args)

        with patch.object(conv.asyncio, "to_thread", fake_to_thread):
            resp = await _upload(
                app, headers, _pdf_bytes("hi"), "a.pdf", "application/pdf"
            )
            get = AsyncMock(return_value=(b"x", {}))
            with patch.object(conv, "s3_get_file", get):
                await conv._load_attached_file(f"conversations/{CONV_ID}/files/u.txt")
        assert resp.status_code == 200
        assert calls == [conv.extract_text, conv.extract_text]

    async def test_other_users_conversation_is_404(self, app, headers, s3):
        resp = await _upload(
            app, headers, b"hi", "a.txt", "text/plain", conversation_id=OTHER_ID
        )
        assert resp.status_code == 404
        s3.upload_file.assert_not_awaited()

    async def test_requires_auth(self, app, s3):
        resp = await _upload(app, {}, b"hi", "a.txt", "text/plain")
        assert resp.status_code == 401


class TestStreamQueryFileKey:
    async def test_rejects_file_key_from_other_conversation(self, app, headers, s3):
        from blueprints.conversation import runs

        add = AsyncMock()
        owned = AsyncMock(return_value=SimpleNamespace(id=CONV_ID))
        with (
            patch.object(runs.Conversation, "get_or_none", owned),
            patch.object(runs.ChatRun, "create", AsyncMock()) as create,
            patch.object(runs, "add_message_to_conversation", add),
        ):
            resp = await app.test_client().post(
                "/api/conversation/stream-query",
                headers=headers,
                json={
                    "query": "summarise",
                    "conversation_id": CONV_ID,
                    "file_key": f"conversations/{OTHER_ID}/files/x.pdf",
                },
            )
        assert resp.status_code == 400
        create.assert_not_awaited()
        add.assert_not_awaited()

    async def test_load_attached_file_recovers_name_and_text(self):
        get = AsyncMock(return_value=(b"a,b\n1,2", {"filename": "M%C3%A9nage.csv"}))
        with patch.object(conv, "s3_get_file", get):
            name, text = await conv._load_attached_file(
                f"conversations/{CONV_ID}/files/u.csv"
            )
        assert name == "Ménage.csv"
        assert text == "a,b\n1,2"

    async def test_load_attached_file_falls_back_to_key_name(self):
        get = AsyncMock(return_value=(b"x", {}))
        with patch.object(conv, "s3_get_file", get):
            name, _ = await conv._load_attached_file(
                f"conversations/{CONV_ID}/files/u.txt"
            )
        assert name == "u.txt"