diff --git a/pytest.ini b/pytest.ini index d3706950..df485620 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,4 +1,6 @@ [pytest] testpaths=tests markers = - filename(name): specify the filename for the pdf_document fixture \ No newline at end of file + filename(name): specify the filename for the pdf_document fixture +filterwarnings = + ignore::Warning \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py index 8949dee5..2e6a308a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -48,12 +48,9 @@ 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) @@ -61,10 +58,13 @@ def pdf_document(request, layout_model, recognition_model, detection_model) -> D 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