wehub-resource-sync
da6a09ff09
Auto Bump and Release / auto-bump-and-release (push) Has been cancelled
style-check / pre-commit (push) Has been cancelled
unit-test / unit testing with python (push) Has been cancelled
unit-test / unit testing with python 3.10 (push) Has been cancelled
unit-test / unit testing with python 3.11 (push) Has been cancelled
52 lines
1.4 KiB
Python
52 lines
1.4 KiB
Python
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from kotaemon.loaders import MathpixPDFReader, OCRReader, PandasExcelReader
|
|
|
|
from .conftest import skip_when_unstructured_pdf_not_installed
|
|
|
|
input_file = Path(__file__).parent / "resources" / "table.pdf"
|
|
input_file_excel = Path(__file__).parent / "resources" / "dummy.xlsx"
|
|
|
|
|
|
@pytest.fixture
|
|
def fullocr_output():
|
|
with open(
|
|
Path(__file__).parent / "resources" / "fullocr_sample_output.json",
|
|
encoding="utf-8",
|
|
) as f:
|
|
fullocr = json.load(f)
|
|
return fullocr
|
|
|
|
|
|
@pytest.fixture
|
|
def mathpix_output():
|
|
with open(Path(__file__).parent / "resources" / "policy.md", encoding="utf-8") as f:
|
|
content = f.read()
|
|
return content
|
|
|
|
|
|
@skip_when_unstructured_pdf_not_installed
|
|
def test_ocr_reader(fullocr_output):
|
|
reader = OCRReader()
|
|
documents = reader.load_data(input_file, response_content=fullocr_output)
|
|
table_docs = [doc for doc in documents if doc.metadata.get("type", "") == "table"]
|
|
assert len(table_docs) == 2
|
|
|
|
|
|
def test_mathpix_reader(mathpix_output):
|
|
reader = MathpixPDFReader()
|
|
documents = reader.load_data(input_file, response_content=mathpix_output)
|
|
table_docs = [doc for doc in documents if doc.metadata.get("type", "") == "table"]
|
|
assert len(table_docs) == 4
|
|
|
|
|
|
def test_excel_reader():
|
|
reader = PandasExcelReader()
|
|
documents = reader.load_data(
|
|
input_file_excel,
|
|
)
|
|
assert len(documents) == 1
|