Update tests

Author: BenConstable9

image_processing/src/image_processing/layout_holders.py

@@ -6,7 +6,6 @@
 
 
 class FigureHolder(BaseModel):
-
     """A class to hold the figure extracted from the document."""
 
     figure_id: str = Field(..., alias="FigureId")

image_processing/src/image_processing/mark_up_cleaner.py

@@ -18,7 +18,7 @@ def get_sections(self, text) -> list:
             list: The sections related to text
         """
         # Updated regex pattern to capture markdown headers like ### Header
-        combined_pattern = r"^[#]+\s*(.*?)(?=\n|$)"
+        combined_pattern = r"^\s*[#]+\s*(.*?)(?=\n|$)"
         doc_metadata = re.findall(combined_pattern, text, re.MULTILINE)
         return self.clean_sections(doc_metadata)
 
@@ -61,12 +61,14 @@ def remove_markdown_tags(self, text: str, tag_patterns: dict) -> str:
             for tag, pattern in tag_patterns.items():
                 try:
                     # Replace the tags using the specific pattern, keeping the content inside the tags
-                    if tag == "header":
+                    if tag in ["header", "figure"]:
                         text = re.sub(
                             pattern, r"\2", text, flags=re.DOTALL | re.MULTILINE
                         )
                     else:
-                        text = re.sub(pattern, r"\1", text, flags=re.DOTALL)
+                        text = re.sub(
+                            pattern, r"\1", text, flags=re.DOTALL | re.MULTILINE
+                        )
                 except re.error as e:
                     logging.error(f"Regex error for tag '{tag}': {e}")
         except Exception as e:
@@ -110,7 +112,7 @@ def clean_text_and_extract_metadata(
             # Define specific patterns for each tag
             tag_patterns = {
                 "figurecontent": r"<!-- FigureContent=(.*?)-->",
-                "figure": r"<figure(?:\s+FigureId=\"[^\"]*\")?>(.*?)</figure>",
+                "figure": r"<figure(?:\s+FigureId=(\"[^\"]*\"|'[^']*'))?>(.*?)</figure>",
                 "figures": r"\(figures/\d+\)(.*?)\(figures/\d+\)",
                 "figcaption": r"<figcaption>(.*?)</figcaption>",
                 "header": r"^\s*(#{1,6})\s*(.*?)\s*$",

image_processing/tests/image_processing/test_mark_up_cleaner.py

@@ -61,7 +61,7 @@ def test_remove_markdown_tags(cleaner):
     """
     tag_patterns = {
         "figurecontent": r"<!-- FigureContent=(.*?)-->",
-        "figure": r"<figure(?:\s+FigureId=\"[^\"]*\")?>(.*?)</figure>",
+        "figure": r"<figure(?:\s+FigureId=(\"[^\"]*\"|'[^']*'))?>(.*?)</figure>",
     }
     cleaned_text = cleaner.remove_markdown_tags(text, tag_patterns)
     assert "Some figure" in cleaned_text
@@ -77,12 +77,18 @@ def test_clean_text_and_extract_metadata(cleaner, sample_text, figures):
     assert result["chunk_sections"] == ["Header 1", "Header 2"]
     assert result["chunk_figures"] == [
         {
-            "figure_id": "12345",
-            "uri": "https://example.com/12345.png",
-            "description": "Figure 1",
+            "FigureId": "12345",
+            "Caption": None,
+            "offset": 0,
+            "length": 8,
+            "PageNumber": None,
+            "Uri": "https://example.com/12345.png",
+            "Description": "Figure 1",
+            "Data": None,
         }
     ]
     assert "chunk_cleaned" in result
+    print(result["chunk_cleaned"])
     assert "FigureId='12345'" not in result["chunk_cleaned"]
 
 
@@ -97,11 +103,15 @@ async def test_clean(cleaner, sample_text, figures):
                     "figure_id": "12345",
                     "uri": "https://example.com/12345.png",
                     "description": "Figure 1",
+                    "offset": 0,
+                    "length": 8,
                 },
                 {
                     "figure_id": "123456789",
                     "uri": "https://example.com/123456789.png",
                     "description": "Figure 2",
+                    "offset": 10,
+                    "length": 8,
                 },
             ],
         },
@@ -112,3 +122,8 @@ async def test_clean(cleaner, sample_text, figures):
     assert result["data"] is not None
     assert result["data"]["chunk_cleaned"]
     assert "errors" not in result or result["errors"] is None
+    assert "chunk_mark_up" in result["data"]
+    assert "chunk_sections" in result["data"]
+    assert "chunk_figures" in result["data"]
+    assert len(result["data"]["chunk_figures"]) == 1
+    assert result["data"]["chunk_figures"][0]["FigureId"] == "12345"