update tests

Author: qued

test_unstructured_inference/inference/test_layout.py

@@ -13,6 +13,7 @@
     EmbeddedTextRegion,
     ImageTextRegion,
 )
+from unstructured_inference.constants import IsExtracted
 from unstructured_inference.models.unstructuredmodel import (
     UnstructuredElementExtractionModel,
     UnstructuredObjectDetectionModel,
@@ -34,7 +35,7 @@ def mock_initial_layout():
         6,
         8,
         text="A very repetitive narrative. " * 10,
-        is_extracted=True,
+        is_extracted=IsExtracted.TRUE,
     )
 
     title_block = EmbeddedTextRegion.from_coords(
@@ -43,7 +44,7 @@ def mock_initial_layout():
         3,
         4,
         text="A Catchy Title",
-        is_extracted=True,
+        is_extracted=IsExtracted.TRUE,
     )
 
     return [text_block, title_block]

test_unstructured_inference/inference/test_layout_element.py

@@ -1,5 +1,5 @@
 from unstructured_inference.inference.layoutelement import LayoutElement, TextRegion
-from unstructured_inference.constants import Source
+from unstructured_inference.constants import IsExtracted, Source
 
 
 def test_layout_element_to_dict(mock_layout_element):
@@ -28,20 +28,22 @@ def test_layoutelement_inheritance_works_correctly():
 
     # Create a TextRegion with both source and text_source
     region = TextRegion.from_coords(
-        0, 0, 10, 10, text="test", source=Source.YOLOX, is_extracted=True
+        0, 0, 10, 10, text="test", source=Source.YOLOX, is_extracted=IsExtracted.TRUE
     )
 
     # Convert to LayoutElement
     element = LayoutElement.from_region(region)
 
     # Check that both properties are preserved
     assert element.source == Source.YOLOX, "LayoutElement should inherit source from TextRegion"
-    assert element.is_extracted, "LayoutElement should inherit is_extracted from TextRegion"
+    assert (
+        element.is_extracted == IsExtracted.TRUE
+    ), "LayoutElement should inherit is_extracted from TextRegion"
 
     # Check that to_dict() works correctly
     d = element.to_dict()
     assert d["source"] == Source.YOLOX
-    assert d["is_extracted"]
+    assert d["is_extracted"] == IsExtracted.TRUE
 
     # Check that we can set source directly on LayoutElement
     element.source = Source.DETECTRON2_ONNX

test_unstructured_inference/test_elements.py

@@ -5,7 +5,7 @@
 import numpy as np
 import pytest
 
-from unstructured_inference.constants import Source
+from unstructured_inference.constants import IsExtracted, Source
 from unstructured_inference.inference import elements
 from unstructured_inference.inference.elements import (
     Rectangle,
@@ -456,8 +456,8 @@ def test_layoutelements_concatenate():
                 ]
             ),
             texts=np.array(["0", "1", "2", "3", "4"]),
-            is_extracted_array=np.array([True] * 5),
-            is_extracted=True,
+            is_extracted_array=np.array([IsExtracted.TRUE] * 5),
+            is_extracted=IsExtracted.TRUE,
         ),
         LayoutElements(
             element_coords=np.array(
@@ -473,7 +473,7 @@ def test_layoutelements_concatenate():
             sources=np.array([Source.YOLOX] * 5),
             source=Source.YOLOX,
             is_extracted_array=np.array([] * 5),
-            is_extracted=True,
+            is_extracted=IsExtracted.TRUE,
             element_probs=np.array([0.0, 0.1, 0.2, 0.3, 0.4]),
         ),
     ],
@@ -495,15 +495,17 @@ def test_textregions_from_list_collects_sources():
     from unstructured_inference.inference.elements import TextRegion
 
     regions = [
-        TextRegion.from_coords(0, 0, 10, 10, text="first", source=Source.YOLOX, is_extracted=True),
+        TextRegion.from_coords(
+            0, 0, 10, 10, text="first", source=Source.YOLOX, is_extracted=IsExtracted.TRUE
+        ),
         TextRegion.from_coords(
             10,
             10,
             20,
             20,
             text="second",
             source=Source.DETECTRON2_ONNX,
-            is_extracted=True,
+            is_extracted=IsExtracted.TRUE,
         ),
     ]
 
@@ -529,7 +531,9 @@ def test_textregions_iter_elements_preserves_source():
     from unstructured_inference.inference.elements import TextRegion
 
     regions = [
-        TextRegion.from_coords(0, 0, 10, 10, text="first", source=Source.YOLOX, is_extracted=True),
+        TextRegion.from_coords(
+            0, 0, 10, 10, text="first", source=Source.YOLOX, is_extracted=IsExtracted.TRUE
+        ),
     ]
     text_regions = TextRegions.from_list(regions)
 
@@ -544,15 +548,17 @@ def test_textregions_slice_preserves_sources():
     from unstructured_inference.inference.elements import TextRegion
 
     regions = [
-        TextRegion.from_coords(0, 0, 10, 10, text="first", source=Source.YOLOX, is_extracted=True),
+        TextRegion.from_coords(
+            0, 0, 10, 10, text="first", source=Source.YOLOX, is_extracted=IsExtracted.TRUE
+        ),
         TextRegion.from_coords(
             10,
             10,
             20,
             20,
             text="second",
             source=Source.DETECTRON2_ONNX,
-            is_extracted=True,
+            is_extracted=IsExtracted.TRUE,
         ),
     ]
     text_regions = TextRegions.from_list(regions)
@@ -581,7 +587,7 @@ def test_textregions_from_coords_accepts_source():
     """Test that TextRegion.from_coords() accepts source parameter"""
     # This should fail because from_coords() doesn't accept source parameter
     region = TextRegion.from_coords(
-        0, 0, 10, 10, text="test", source=Source.YOLOX, is_extracted=True
+        0, 0, 10, 10, text="test", source=Source.YOLOX, is_extracted=IsExtracted.TRUE
     )
 
     assert region.source == Source.YOLOX