PR improvements: some refactoring, better test verification

Author: Val Brodsky

labelbox/data/annotation_types/ner/document_entity.py

@@ -19,5 +19,4 @@ def validate_page(cls, v):
 
 class DocumentEntity(_CamelCaseMixin, BaseModel):
     """ Represents a text entity """
-    name: str
     text_selections: List[DocumentTextSelection]

labelbox/data/serialization/ndjson/objects.py

@@ -505,8 +505,8 @@ def lookup_object(
 # NOTE: Deserialization of subclasses in pydantic is a known PIA, see here https://blog.devgenius.io/deserialize-child-classes-with-pydantic-that-gonna-work-784230e1cf83
 # I could implement the registry approach suggested there, but I found that if I list subclass (that has more attributes) before the parent class, it works
 # This is a bit of a hack, but it works for now
-NERTextType = Union[NDConversationEntity, NDTextEntity]
+NDEntityType = Union[NDConversationEntity, NDTextEntity]
 NDObjectType = Union[NDLine, NDPolygon, NDPoint, NDRectangle, NDMask,
-                     NERTextType, NDDocumentEntity]
+                     NDEntityType, NDDocumentEntity]
 
 NDFrameObjectType = NDFrameRectangle, NDFramePoint, NDFrameLine

tests/data/annotation_types/test_ner.py

@@ -11,15 +11,10 @@ def test_ner():
 
 
 def test_document_entity():
-    document_entity = DocumentEntity(name="tool_name",
-                                     text_selections=[
-                                         DocumentTextSelection(
-                                             token_ids=["1", "2"],
-                                             group_id="1",
-                                             page=1)
-                                     ])
-
-    assert document_entity.name == "tool_name"
+    document_entity = DocumentEntity(text_selections=[
+        DocumentTextSelection(token_ids=["1", "2"], group_id="1", page=1)
+    ])
+
     assert document_entity.text_selections[0].token_ids == ["1", "2"]
     assert document_entity.text_selections[0].group_id == "1"
     assert document_entity.text_selections[0].page == 1

tests/integration/annotation_import/test_bulk_import_request.py

@@ -311,7 +311,7 @@ def test_pdf_document_entity(client, configured_project_without_data_rows,
         page=1)
 
     entities_annotation_document_entity = DocumentEntity(
-        name="named_entity", text_selections=[document_text_selection])
+        text_selections=[document_text_selection])
     entities_annotation = ObjectAnnotation(
         name="named-entity", value=entities_annotation_document_entity)
 

tests/integration/annotation_import/test_conversation_import.py

@@ -10,8 +10,7 @@
 def test_conversation_entity(client, configured_project_without_data_rows,
                              dataset_conversation_entity, rand_gen):
 
-    conversation_entity_annotation = ConversationEntity(name="named-entity",
-                                                        start=0,
+    conversation_entity_annotation = ConversationEntity(start=0,
                                                         end=8,
                                                         message_id="4")
 
@@ -41,3 +40,15 @@ def test_conversation_entity(client, configured_project_without_data_rows,
     import_annotations.wait_until_done()
 
     assert import_annotations.errors == []
+
+    exported_labels = configured_project_without_data_rows.label_generator()
+    for label in exported_labels:
+        assert len(
+            label.annotations) == 1  # we have created only 1 annotation above
+        annotation = label.annotations[0]
+
+        assert type(annotation) is ConversationEntity
+        assert annotation.name == "named-entity"
+        assert annotation.value.message_id == "4"
+        assert annotation.value.start == 0
+        assert annotation.value.end == 8