Spaces:
Runtime error
Runtime error
File size: 5,426 Bytes
b13c030 49206fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | import pytest
from uuid import uuid4
from unittest.mock import patch
from app.models import Document, Chunk
def test_upload_and_process_pdf(client, db_session):
client.post("/api/v1/auth/signup", json={"email": "doc@test.com", "password": "Password123"})
token = client.post("/api/v1/auth/login", json={"email": "doc@test.com", "password": "Password123"}).json()["access_token"]
headers = {"Authorization": f"Bearer {token}"}
key = str(uuid4())
hub_res = client.post("/api/v1/hubs", headers={**headers, "X-Idempotency-Key": key}, json={"name": "PDF Hub"})
hub_id = hub_res.json()["id"]
file_content = b"This is a test document. It has multiple sentences to chunk."
upload_key = str(uuid4())
files = {"file": ("test.txt", file_content, "text/plain")}
upload_res = client.post(f"/api/v1/hubs/{hub_id}/documents", headers={**headers, "X-Idempotency-Key": upload_key}, files=files)
assert upload_res.status_code == 202
doc_id = upload_res.json()["id"]
assert upload_res.json()["status"] == "pending"
with patch("app.routers.documents.embed_texts", return_value=[[0.1]*768]):
res1 = client.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers=headers)
assert res1.json()["status"] == "chunking"
res2 = client.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers=headers)
assert res2.json()["status"] == "embedding"
res3 = client.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers=headers)
assert res3.json()["status"] == "ready"
assert db_session.query(Chunk).count() > 0
def test_process_is_idempotent(client, db_session):
client.post("/api/v1/auth/signup", json={"email": "idem_doc@test.com", "password": "Password123"})
token = client.post("/api/v1/auth/login", json={"email": "idem_doc@test.com", "password": "Password123"}).json()["access_token"]
headers = {"Authorization": f"Bearer {token}"}
hub_res = client.post("/api/v1/hubs", headers={**headers, "X-Idempotency-Key": str(uuid4())}, json={"name": "Idem Doc Hub"})
hub_id = hub_res.json()["id"]
files = {"file": ("test.txt", b"Test text for idempotency.", "text/plain")}
upload_res = client.post(f"/api/v1/hubs/{hub_id}/documents", headers={**headers, "X-Idempotency-Key": str(uuid4())}, files=files)
doc_id = upload_res.json()["id"]
key1 = str(uuid4())
with patch("app.routers.documents.embed_texts", return_value=[[0.1]*768]):
res1 = client.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers={**headers, "X-Idempotency-Key": key1})
assert res1.json()["status"] == "chunking"
res1_retry = client.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers={**headers, "X-Idempotency-Key": key1})
assert res1_retry.json()["status"] == "chunking"
key2 = str(uuid4())
res2 = client.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers={**headers, "X-Idempotency-Key": key2})
assert res2.json()["status"] == "embedding"
def test_upload_unsupported_mime_type(client):
client.post("/api/v1/auth/signup", json={"email": "mime@test.com", "password": "Password123"})
token = client.post("/api/v1/auth/login", json={"email": "mime@test.com", "password": "Password123"}).json()["access_token"]
headers = {"Authorization": f"Bearer {token}"}
hub_res = client.post("/api/v1/hubs", headers={**headers, "X-Idempotency-Key": str(uuid4())}, json={"name": "Mime Hub"})
hub_id = hub_res.json()["id"]
files = {"file": ("test.png", b"fake image data", "image/png")}
res = client.post(f"/api/v1/hubs/{hub_id}/documents", headers={**headers, "X-Idempotency-Key": str(uuid4())}, files=files)
assert res.status_code == 415
assert res.json()["error"]["code"] == "UNSUPPORTED_MIME_TYPE"
def test_upload_file_too_large(client):
client.post("/api/v1/auth/signup", json={"email": "large@test.com", "password": "Password123"})
token = client.post("/api/v1/auth/login", json={"email": "large@test.com", "password": "Password123"}).json()["access_token"]
headers = {"Authorization": f"Bearer {token}"}
hub_res = client.post("/api/v1/hubs", headers={**headers, "X-Idempotency-Key": str(uuid4())}, json={"name": "Large Hub"})
hub_id = hub_res.json()["id"]
# 11MB file
large_content = b"a" * (11 * 1024 * 1024)
files = {"file": ("large.txt", large_content, "text/plain")}
res = client.post(f"/api/v1/hubs/{hub_id}/documents", headers={**headers, "X-Idempotency-Key": str(uuid4())}, files=files)
assert res.status_code == 413
assert res.json()["error"]["code"] == "DOCUMENT_TOO_LARGE"
def test_process_document_not_found(client):
client.post("/api/v1/auth/signup", json={"email": "notfound@test.com", "password": "Password123"})
token = client.post("/api/v1/auth/login", json={"email": "notfound@test.com", "password": "Password123"}).json()["access_token"]
headers = {"Authorization": f"Bearer {token}"}
hub_res = client.post("/api/v1/hubs", headers={**headers, "X-Idempotency-Key": str(uuid4())}, json={"name": "NF Hub"})
hub_id = hub_res.json()["id"]
random_doc_id = uuid4()
res = client.post(f"/api/v1/hubs/{hub_id}/documents/{random_doc_id}/process", headers=headers)
assert res.status_code == 404
assert res.json()["error"]["code"] == "DOCUMENT_NOT_FOUND" |