346 lines
12 KiB
Python
346 lines
12 KiB
Python
"""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"
|
|
|
|
|
|
class TestImageOwnership:
|
|
async def test_presigns_only_keys_in_owned_conversations(self, app, headers, s3):
|
|
presign = AsyncMock(return_value="https://s3.example/signed")
|
|
with patch.object(conv, "s3_presigned_url", presign):
|
|
client = app.test_client()
|
|
own = await client.get(
|
|
f"/api/conversation/image/conversations/{CONV_ID}/a.jpg",
|
|
headers=headers,
|
|
)
|
|
assert own.status_code == 200
|
|
for key in (
|
|
f"conversations/{OTHER_ID}/a.jpg",
|
|
f"conversations/{CONV_ID}/../{OTHER_ID}/a.jpg",
|
|
"conversations/not-a-uuid/a.jpg",
|
|
"beli/profile.jpg",
|
|
):
|
|
resp = await client.get(
|
|
f"/api/conversation/image/{key}", headers=headers
|
|
)
|
|
assert resp.status_code == 404, key
|
|
presign.assert_awaited_once_with(f"conversations/{CONV_ID}/a.jpg")
|
|
|
|
async def test_image_upload_requires_owned_conversation(self, app, headers, s3):
|
|
with patch.object(conv.blueprints.users.models.User, "get", AsyncMock()):
|
|
resp = await self._upload_image(app, headers)
|
|
assert resp.status_code == 404
|
|
s3.upload_image.assert_not_awaited()
|
|
|
|
@staticmethod
|
|
async def _upload_image(app, headers):
|
|
return await app.test_client().post(
|
|
"/api/conversation/upload-image",
|
|
headers=headers,
|
|
form={"conversation_id": OTHER_ID},
|
|
files={
|
|
"file": FileStorage(
|
|
io.BytesIO(_png_bytes()), filename="a.png", content_type="image/png"
|
|
)
|
|
},
|
|
)
|