Add default value of None for optional parameters, and add more tests

google/cloud/documentai_toolbox/wrappers/entity.py

@@ -57,13 +57,13 @@ class Entity:
     page_offset: dataclasses.InitVar[Optional[int]] = 0
 
     type_: str = dataclasses.field(init=False)
-    mention_text: Optional[str] = dataclasses.field(init=False, default="")
-    normalized_text: Optional[str] = dataclasses.field(init=False, default="")
+    mention_text: Optional[str] = dataclasses.field(init=False, default=None)
+    normalized_text: Optional[str] = dataclasses.field(init=False, default=None)
 
     # Not populated in Classifier output
-    start_page: Optional[int] = dataclasses.field(init=False)
+    start_page: Optional[int] = dataclasses.field(init=False, default=None)
     # Only Populated for Splitter Output
-    end_page: Optional[int] = dataclasses.field(init=False)
+    end_page: Optional[int] = dataclasses.field(init=False, default=None)
 
     _image: Optional[Image.Image] = dataclasses.field(init=False, default=None)
 

tests/unit/test_document.py

@@ -226,6 +226,15 @@ def test_entities_from_shards_classifier(get_bytes_classifier_mock):
     assert actual[0].type_ == "computer_vision"
     assert round(actual[0].documentai_object.confidence, 8) == 0.47925246
     assert actual[0].documentai_object.id == "0"
+    assert actual[1].type_ == "crypto"
+    assert round(actual[1].documentai_object.confidence, 8) == 0.0433604
+    assert actual[1].documentai_object.id == "1"
+    assert actual[2].type_ == "med_tech"
+    assert round(actual[2].documentai_object.confidence, 8) == 0.26732057
+    assert actual[2].documentai_object.id == "2"
+    assert actual[3].type_ == "other"
+    assert round(actual[3].documentai_object.confidence, 8) == 0.2100666
+    assert actual[3].documentai_object.id == "3"
 
 
 @mock.patch("google.cloud.documentai_toolbox.wrappers.document.documentai")