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
Files changed (2) hide show
  1. pytest.ini +3 -1
  2. 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 pdf_document(request, layout_model, recognition_model, detection_model) -> Document:
52
  mark = request.node.get_closest_marker("filename")
53
- if mark is None:
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
- provider = PdfProvider(temp_pdf.name)
 
 
66
  layout_builder = LayoutBuilder(layout_model)
67
  ocr_builder = OcrBuilder(detection_model, recognition_model)
68
  builder = DocumentBuilder()
69
- document = builder(provider, layout_builder, ocr_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