Spaces:
Running
Running
File size: 2,056 Bytes
d825c91 |
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 |
import pytest
from fastapi import UploadFile
from io import BytesIO
from services.document_service import document_service
import os
@pytest.fixture
def sample_pdf():
return BytesIO(b"%PDF-1.4\n%Test PDF content")
@pytest.fixture
def sample_docx():
return BytesIO(b"PK\x03\x04\x14\x00\x00\x00\x08\x00")
@pytest.fixture
def sample_image():
return BytesIO(b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR")
@pytest.mark.asyncio
async def test_save_upload_file(sample_pdf):
file = UploadFile(
filename="test.pdf",
file=sample_pdf
)
file_path = await document_service.save_upload_file(file)
assert file_path.endswith("test.pdf")
assert os.path.exists(file_path)
@pytest.mark.asyncio
async def test_process_pdf(sample_pdf):
file = UploadFile(
filename="test.pdf",
file=sample_pdf
)
file_path = await document_service.save_upload_file(file)
result = await document_service.process_document(file_path)
assert "text" in result
assert result["type"] == "pdf"
@pytest.mark.asyncio
async def test_process_docx(sample_docx):
file = UploadFile(
filename="test.docx",
file=sample_docx
)
file_path = await document_service.save_upload_file(file)
result = await document_service.process_document(file_path)
assert "text" in result
assert result["type"] == "word"
@pytest.mark.asyncio
async def test_process_image(sample_image):
file = UploadFile(
filename="test.png",
file=sample_image
)
file_path = await document_service.save_upload_file(file)
result = await document_service.process_document(file_path)
assert "text" in result
assert result["type"] == "image"
@pytest.mark.asyncio
async def test_segment_document():
text = "Line 1\nLine 2\nLine 3\nLine 4\nLine 5"
segments = await document_service.segment_document(text, max_segment_size=20)
assert len(segments) > 0
assert all(len(segment) <= 20 for segment in segments) |