"""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('\n') assert block.endswith(f"{TRUNCATION_MARKER}\n") 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"