Moses Paul R
commited on
Commit
·
5380a39
1
Parent(s):
10e6ee7
update pytest ini and conftest to ignore warnings and fix pdf provider [skip ci]
Browse files- pytest.ini +3 -1
- tests/conftest.py +7 -7
pytest.ini
CHANGED
|
@@ -1,4 +1,6 @@
|
|
| 1 |
[pytest]
|
| 2 |
testpaths=tests
|
| 3 |
markers =
|
| 4 |
-
filename(name): specify the filename for the pdf_document fixture
|
|
|
|
|
|
|
|
|
| 1 |
[pytest]
|
| 2 |
testpaths=tests
|
| 3 |
markers =
|
| 4 |
+
filename(name): specify the filename for the pdf_document fixture
|
| 5 |
+
filterwarnings =
|
| 6 |
+
ignore::Warning
|
tests/conftest.py
CHANGED
|
@@ -48,12 +48,9 @@ def table_rec_model():
|
|
| 48 |
|
| 49 |
|
| 50 |
@pytest.fixture(scope="function")
|
| 51 |
-
def
|
| 52 |
mark = request.node.get_closest_marker("filename")
|
| 53 |
-
if mark
|
| 54 |
-
filename = "adversarial.pdf"
|
| 55 |
-
else:
|
| 56 |
-
filename = mark.args[0]
|
| 57 |
|
| 58 |
dataset = datasets.load_dataset("datalab-to/pdfs", split="train")
|
| 59 |
idx = dataset['filename'].index(filename)
|
|
@@ -61,10 +58,13 @@ def pdf_document(request, layout_model, recognition_model, detection_model) -> D
|
|
| 61 |
temp_pdf = tempfile.NamedTemporaryFile(suffix=".pdf")
|
| 62 |
temp_pdf.write(dataset['pdf'][idx])
|
| 63 |
temp_pdf.flush()
|
|
|
|
| 64 |
|
| 65 |
-
|
|
|
|
|
|
|
| 66 |
layout_builder = LayoutBuilder(layout_model)
|
| 67 |
ocr_builder = OcrBuilder(detection_model, recognition_model)
|
| 68 |
builder = DocumentBuilder()
|
| 69 |
-
document = builder(
|
| 70 |
return document
|
|
|
|
| 48 |
|
| 49 |
|
| 50 |
@pytest.fixture(scope="function")
|
| 51 |
+
def pdf_provider(request):
|
| 52 |
mark = request.node.get_closest_marker("filename")
|
| 53 |
+
filename = mark.args[0] if mark else "adversarial.pdf"
|
|
|
|
|
|
|
|
|
|
| 54 |
|
| 55 |
dataset = datasets.load_dataset("datalab-to/pdfs", split="train")
|
| 56 |
idx = dataset['filename'].index(filename)
|
|
|
|
| 58 |
temp_pdf = tempfile.NamedTemporaryFile(suffix=".pdf")
|
| 59 |
temp_pdf.write(dataset['pdf'][idx])
|
| 60 |
temp_pdf.flush()
|
| 61 |
+
yield PdfProvider(temp_pdf.name)
|
| 62 |
|
| 63 |
+
|
| 64 |
+
@pytest.fixture(scope="function")
|
| 65 |
+
def pdf_document(pdf_provider, layout_model, recognition_model, detection_model) -> Document:
|
| 66 |
layout_builder = LayoutBuilder(layout_model)
|
| 67 |
ocr_builder = OcrBuilder(detection_model, recognition_model)
|
| 68 |
builder = DocumentBuilder()
|
| 69 |
+
document = builder(pdf_provider, layout_builder, ocr_builder)
|
| 70 |
return document
|