Add tests

Author: kvilon

libs/labelbox/tests/data/annotation_import/test_relationships.py

@@ -409,3 +409,62 @@ def test_relationship_readonly_explicit_true():
             ),
         )
     assert relationship.value.readonly is True
+
+
+def test_relationship_source_ontology_name():
+    """Test that relationship can be created with source_ontology_name instead of source."""
+    target = ObjectAnnotation(
+        name="e2",
+        value=TextEntity(start=30, end=35),
+    )
+
+    relationship = RelationshipAnnotation(
+        name="rel",
+        value=Relationship(
+            source_ontology_name="test_source",
+            target=target,
+            type=Relationship.Type.UNIDIRECTIONAL,
+        ),
+    )
+    assert relationship.value.source_ontology_name == "test_source"
+    assert relationship.value.source is None
+
+
+def test_relationship_missing_source_validation():
+    """Test that relationship requires either source or source_ontology_name."""
+    target = ObjectAnnotation(
+        name="e2",
+        value=TextEntity(start=30, end=35),
+    )
+
+    with pytest.raises(ValueError, match="Either source or source_ontology_name must be provided"):
+        RelationshipAnnotation(
+            name="rel",
+            value=Relationship(
+                target=target,
+                type=Relationship.Type.UNIDIRECTIONAL,
+            ),
+        )
+
+
+def test_relationship_both_sources_validation():
+    """Test that relationship cannot have both source and source_ontology_name."""
+    source = ObjectAnnotation(
+        name="e1",
+        value=TextEntity(start=10, end=12),
+    )
+    target = ObjectAnnotation(
+        name="e2",
+        value=TextEntity(start=30, end=35),
+    )
+
+    with pytest.raises(ValueError, match="Only one of 'source' or 'source_ontology_name' may be provided"):
+        RelationshipAnnotation(
+            name="rel",
+            value=Relationship(
+                source=source,
+                source_ontology_name="test_source",
+                target=target,
+                type=Relationship.Type.UNIDIRECTIONAL,
+            ),
+        )

libs/labelbox/tests/data/serialization/ndjson/test_relationship.py

@@ -5,6 +5,9 @@
     RelationshipAnnotation,
     Relationship,
     TextEntity,
+    DocumentRectangle,
+    RectangleUnit,
+    Point,
 )
 
 
@@ -285,3 +288,133 @@ def test_readonly_relationships():
         non_readonly_rel_serialized["relationship"]["type"] == "bidirectional"
     )
     assert non_readonly_rel_serialized["relationship"]["readonly"] is False
+
+
+def test_source_ontology_name_relationship():
+    ner_source = ObjectAnnotation(
+        name="e1",
+        value=TextEntity(start=10, end=12),
+    )
+
+    ner_target = ObjectAnnotation(
+        name="e2",
+        value=TextEntity(start=30, end=35),
+    )
+
+    # Test relationship with source
+    regular_relationship = RelationshipAnnotation(
+        name="regular_rel",
+        value=Relationship(
+            source=ner_source,
+            target=ner_target,
+            type=Relationship.Type.UNIDIRECTIONAL,
+        ),
+    )
+
+    # Test relationship with source_ontology_name for PDF target
+    pdf_target = ObjectAnnotation(
+        name="pdf_region",
+        value=DocumentRectangle(
+            start=Point(x=0.5, y=0.5),
+            end=Point(x=0.7, y=0.7),
+            page=1,
+            unit=RectangleUnit.PERCENT,
+        ),
+    )
+
+    ontology_relationship = RelationshipAnnotation(
+        name="ontology_rel",
+        value=Relationship(
+            source_ontology_name="Person",
+            target=pdf_target,
+            type=Relationship.Type.UNIDIRECTIONAL,
+        ),
+    )
+
+    label = Label(
+        data={"uid": "clqbkpy236syk07978v3pscw1"},
+        annotations=[
+            ner_source,
+            ner_target,
+            pdf_target,
+            regular_relationship,
+            ontology_relationship,
+        ],
+    )
+
+    serialized_label = list(NDJsonConverter.serialize([label]))
+
+    ner_source_serialized = next(
+        annotation
+        for annotation in serialized_label
+        if annotation["name"] == ner_source.name
+    )
+    ner_target_serialized = next(
+        annotation
+        for annotation in serialized_label
+        if annotation["name"] == ner_target.name
+    )
+    pdf_target_serialized = next(
+        annotation
+        for annotation in serialized_label
+        if annotation["name"] == pdf_target.name
+    )
+    regular_rel_serialized = next(
+        annotation
+        for annotation in serialized_label
+        if annotation["name"] == regular_relationship.name
+    )
+    ontology_rel_serialized = next(
+        annotation
+        for annotation in serialized_label
+        if annotation["name"] == ontology_relationship.name
+    )
+
+    # Verify regular relationship
+    assert (
+        regular_rel_serialized["relationship"]["source"]
+        == ner_source_serialized["uuid"]
+    )
+    assert (
+        regular_rel_serialized["relationship"]["target"]
+        == ner_target_serialized["uuid"]
+    )
+    assert regular_rel_serialized["relationship"]["type"] == "unidirectional"
+
+    # Verify relationship with source_ontology_name
+    assert (
+        ontology_rel_serialized["relationship"]["sourceOntologyName"] == "Person"
+    )
+    assert (
+        ontology_rel_serialized["relationship"]["target"]
+        == pdf_target_serialized["uuid"]
+    )
+    assert ontology_rel_serialized["relationship"]["type"] == "unidirectional"
+
+    # Test that providing both source and source_ontology_name raises an error
+    try:
+        RelationshipAnnotation(
+            name="invalid_rel",
+            value=Relationship(
+                source=ner_source,
+                source_ontology_name="Person",
+                target=pdf_target,
+                type=Relationship.Type.UNIDIRECTIONAL,
+            ),
+        )
+        assert False, "Expected ValueError for providing both source and source_ontology_name"
+    except Exception as e:
+        assert "Value error, Only one of 'source' or 'source_ontology_name' may be provided" in str(e)
+
+    # Test that providing neither source nor source_ontology_name raises an error
+    try:
+        RelationshipAnnotation(
+            name="invalid_rel",
+            value=Relationship(
+                target=pdf_target,
+                type=Relationship.Type.UNIDIRECTIONAL,
+            ),
+        )
+        assert False, "Expected ValueError for providing neither source nor source_ontology_name"
+    except Exception as e:
+        assert "Value error, Either source or source_ontology_name must be provided" in str(e)