Skip to content

Commit

Permalink
update pytest ini and conftest to ignore warnings and fix pdf provide…
Browse files Browse the repository at this point in the history
…r [skip ci]
  • Loading branch information
iammosespaulr committed Nov 18, 2024
1 parent 028ef50 commit 19d313c
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 8 deletions.
4 changes: 3 additions & 1 deletion pytest.ini
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
[pytest]
testpaths=tests
markers =
filename(name): specify the filename for the pdf_document fixture
filename(name): specify the filename for the pdf_document fixture
filterwarnings =
ignore::Warning
14 changes: 7 additions & 7 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,23 +48,23 @@ def table_rec_model():


@pytest.fixture(scope="function")
def pdf_document(request, layout_model, recognition_model, detection_model) -> Document:
def pdf_provider(request):
mark = request.node.get_closest_marker("filename")
if mark is None:
filename = "adversarial.pdf"
else:
filename = mark.args[0]
filename = mark.args[0] if mark else "adversarial.pdf"

dataset = datasets.load_dataset("datalab-to/pdfs", split="train")
idx = dataset['filename'].index(filename)

temp_pdf = tempfile.NamedTemporaryFile(suffix=".pdf")
temp_pdf.write(dataset['pdf'][idx])
temp_pdf.flush()
yield PdfProvider(temp_pdf.name)

provider = PdfProvider(temp_pdf.name)

@pytest.fixture(scope="function")
def pdf_document(pdf_provider, layout_model, recognition_model, detection_model) -> Document:
layout_builder = LayoutBuilder(layout_model)
ocr_builder = OcrBuilder(detection_model, recognition_model)
builder = DocumentBuilder()
document = builder(provider, layout_builder, ocr_builder)
document = builder(pdf_provider, layout_builder, ocr_builder)
return document

0 comments on commit 19d313c

Please sign in to comment.