curo1305
108 lines
3.3 KiB
Python
108 lines
3.3 KiB
Python
def test_upload_txt_no_classify(client, sample_txt):
|
|
with open(sample_txt, "rb") as f:
|
|
resp = client.post(
|
|
"/api/documents/upload",
|
|
files={"file": ("sample.txt", f, "text/plain")},
|
|
data={"auto_classify": "false"},
|
|
)
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert data["original_name"] == "sample.txt"
|
|
assert "extracted_text" in data
|
|
assert "invoices" in data["extracted_text"].lower() or len(data["extracted_text"]) > 0
|
|
assert data["topics"] == []
|
|
assert "id" in data
|
|
|
|
|
|
def test_upload_pdf_no_classify(client, sample_pdf):
|
|
with open(sample_pdf, "rb") as f:
|
|
resp = client.post(
|
|
"/api/documents/upload",
|
|
files={"file": ("sample.pdf", f, "application/pdf")},
|
|
data={"auto_classify": "false"},
|
|
)
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert data["mime_type"] == "application/pdf"
|
|
assert len(data["extracted_text"]) > 0
|
|
|
|
|
|
def test_list_documents(client, sample_txt):
|
|
with open(sample_txt, "rb") as f:
|
|
client.post(
|
|
"/api/documents/upload",
|
|
files={"file": ("a.txt", f, "text/plain")},
|
|
data={"auto_classify": "false"},
|
|
)
|
|
resp = client.get("/api/documents")
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert data["total"] == 1
|
|
assert len(data["items"]) == 1
|
|
|
|
|
|
def test_list_documents_filter_by_topic(client, sample_txt):
|
|
with open(sample_txt, "rb") as f:
|
|
upload = client.post(
|
|
"/api/documents/upload",
|
|
files={"file": ("a.txt", f, "text/plain")},
|
|
data={"auto_classify": "false"},
|
|
).json()
|
|
|
|
import services.storage as st
|
|
st.update_document_topics(upload["id"], ["finance"])
|
|
|
|
resp = client.get("/api/documents?topic=finance")
|
|
assert resp.json()["total"] == 1
|
|
|
|
resp2 = client.get("/api/documents?topic=legal")
|
|
assert resp2.json()["total"] == 0
|
|
|
|
|
|
def test_get_document(client, sample_txt):
|
|
with open(sample_txt, "rb") as f:
|
|
upload = client.post(
|
|
"/api/documents/upload",
|
|
files={"file": ("a.txt", f, "text/plain")},
|
|
data={"auto_classify": "false"},
|
|
).json()
|
|
|
|
resp = client.get(f"/api/documents/{upload['id']}")
|
|
assert resp.status_code == 200
|
|
assert resp.json()["id"] == upload["id"]
|
|
|
|
|
|
def test_get_document_not_found(client):
|
|
resp = client.get("/api/documents/nonexistent")
|
|
assert resp.status_code == 404
|
|
|
|
|
|
def test_delete_document(client, sample_txt):
|
|
with open(sample_txt, "rb") as f:
|
|
upload = client.post(
|
|
"/api/documents/upload",
|
|
files={"file": ("a.txt", f, "text/plain")},
|
|
data={"auto_classify": "false"},
|
|
).json()
|
|
|
|
resp = client.delete(f"/api/documents/{upload['id']}")
|
|
assert resp.status_code == 200
|
|
assert resp.json()["success"] is True
|
|
|
|
resp2 = client.get(f"/api/documents/{upload['id']}")
|
|
assert resp2.status_code == 404
|
|
|
|
|
|
def test_delete_document_not_found(client):
|
|
resp = client.delete("/api/documents/nonexistent")
|
|
assert resp.status_code == 404
|
|
|
|
|
|
def test_upload_empty_file(client):
|
|
resp = client.post(
|
|
"/api/documents/upload",
|
|
files={"file": ("empty.txt", b"", "text/plain")},
|
|
data={"auto_classify": "false"},
|
|
)
|
|
assert resp.status_code == 400
|