Merge pull request #369 from VikParuchuri/vik_v2

Reorganize tests
VikParuchuri · Nov 19, 2024 · f89089c · f89089c
2 parents 10f6ed0 + b191e17
commit f89089c
Show file tree

Hide file tree

Showing 12 changed files with 46 additions and 17 deletions.
diff --git a/marker/v2/schema/blocks/handwriting.py b/marker/v2/schema/blocks/handwriting.py
@@ -4,3 +4,8 @@
 
 class Handwriting(Block):
     block_type: BlockTypes = BlockTypes.Handwriting
+
+    def assemble_html(self, child_blocks, parent_structure):
+        template = super().assemble_html(child_blocks, parent_structure)
+        template = template.replace("\n", " ")
+        return f"<p>{template}</p>"
diff --git a/tests/test_document_builder.py → tests/builders/test_document_builder.py b/tests/test_document_builder.py → tests/builders/test_document_builder.py
diff --git a/tests/test_garbled_pdf.py → tests/builders/test_garbled_pdf.py b/tests/test_garbled_pdf.py → tests/builders/test_garbled_pdf.py
@@ -3,7 +3,7 @@
 
 
 @pytest.mark.filename("water_damage.pdf")
-def test_ocr_pipeline(pdf_document):
+def test_garbled_pdf(pdf_document):
     assert pdf_document.pages[0].structure[0] == '/page/0/Table/0'
 
     table_block = pdf_document.pages[0].get_block(pdf_document.pages[0].structure[0])

diff --git a/tests/test_ocr_pipeline.py → tests/builders/test_ocr_pipeline.py b/tests/test_ocr_pipeline.py → tests/builders/test_ocr_pipeline.py
diff --git a/tests/test_overriding.py → tests/builders/test_overriding.py b/tests/test_overriding.py → tests/builders/test_overriding.py
diff --git a/tests/test_structure.py → tests/builders/test_structure.py b/tests/test_structure.py → tests/builders/test_structure.py
diff --git a/tests/test_equation_processor.py → tests/processors/test_equation_processor.py b/tests/test_equation_processor.py → tests/processors/test_equation_processor.py
diff --git a/tests/test_table_processor.py → tests/processors/test_table_processor.py b/tests/test_table_processor.py → tests/processors/test_table_processor.py
diff --git a/tests/test_pdf_provider.py → tests/providers/test_pdf_provider.py b/tests/test_pdf_provider.py → tests/providers/test_pdf_provider.py
diff --git a/tests/renderers/test_markdown_renderer.py b/tests/renderers/test_markdown_renderer.py
@@ -0,0 +1,21 @@
+import pytest
+
+from marker.v2.renderers.markdown import MarkdownRenderer
+
+
+@pytest.mark.config({"page_range": [0]})
+def test_markdown_renderer(pdf_document):
+    renderer = MarkdownRenderer()
+    md = renderer(pdf_document).markdown
+
+    # Verify markdown
+    assert '# Subspace Adversarial Training' in md
+
+
+@pytest.mark.config({"page_range": [0, 1], "paginate_output": True})
+def test_markdown_renderer_pagination(pdf_document):
+    renderer = MarkdownRenderer({"paginate_output": True})
+    md = renderer(pdf_document).markdown
+
+    assert "{0}-" in md
+    assert "{1}-" in md
diff --git a/tests/schema/groups/test_list_grouping.py b/tests/schema/groups/test_list_grouping.py
@@ -0,0 +1,19 @@
+import pytest
+
+from marker.v2.builders.structure import StructureBuilder
+from marker.v2.schema import BlockTypes
+
+
+@pytest.mark.config({"page_range": [4]})
+def test_list_grouping(pdf_document):
+    structure = StructureBuilder()
+    structure(pdf_document)
+
+    page = pdf_document.pags[0]
+    list_groups = []
+    for block in page.children:
+        if block.block_type == BlockTypes.ListGroup:
+            list_groups.append(block)
+
+    assert len(list_groups) == 1
+
diff --git a/tests/utils.py b/tests/utils.py
@@ -22,19 +22,3 @@ def setup_pdf_provider(
 
     provider = PdfProvider(temp_pdf.name, config)
     return provider
-
-
-def setup_pdf_document(
-    filename='adversarial.pdf',
-    config=None,
-) -> Document:
-    layout_model = setup_layout_model()
-    recognition_model = setup_recognition_model()
-    detection_model = setup_detection_model()
-
-    provider = setup_pdf_provider(filename, config)
-    layout_builder = LayoutBuilder(layout_model, config)
-    ocr_builder = OcrBuilder(detection_model, recognition_model, config)
-    builder = DocumentBuilder(config)
-    document = builder(provider, layout_builder, ocr_builder)
-    return document