Improved unit tests assertion and updated generator templates

pyatlan/generator/templates/imports.jinja2

@@ -49,6 +49,7 @@ from pyatlan.model.enums import (
     KafkaTopicCompressionType,
     MatillionJobType,
     OpenLineageRunState,
+    SaveSemantic,
     PersonaDomainAction,
     PersonaGlossaryAction,
     PersonaMetadataAction,

pyatlan/generator/templates/methods/asset/asset.jinja2

@@ -39,17 +39,25 @@
         return cls(attributes=cls.Attributes(qualified_name=qualified_name, name=name))
 
     @classmethod
-    def ref_by_guid(cls: type[SelfAsset], guid: str) -> SelfAsset:
+    def ref_by_guid(
+        cls: type[SelfAsset], guid: str, semantic: SaveSemantic = SaveSemantic.REPLACE
+    ) -> SelfAsset:
         retval: SelfAsset = cls(attributes=cls.Attributes())
         retval.guid = guid
+        retval.semantic = semantic
         return retval
 
     @classmethod
-    def ref_by_qualified_name(cls: type[SelfAsset], qualified_name: str) -> SelfAsset:
+    def ref_by_qualified_name(
+        cls: type[SelfAsset],
+        qualified_name: str,
+        semantic: SaveSemantic = SaveSemantic.REPLACE,
+    ) -> SelfAsset:
         ret_value: SelfAsset = cls(
             attributes=cls.Attributes(name="", qualified_name=qualified_name)
         )
         ret_value.unique_attributes = {"qualifiedName": qualified_name}
+        ret_value.semantic = semantic
         return ret_value
 
     @classmethod

pyatlan/generator/templates/referenceable_attributes.jinja2

@@ -176,3 +176,21 @@
     pending_tasks: Optional[list[str]] = Field(None)
 
     unique_attributes: Optional[dict[str, Any]] = Field(None)
+
+    append_relationship_attributes: Optional[dict[str, Any]] = Field(
+        None,
+        alias="appendRelationshipAttributes",
+        description="Map of append relationship attributes.",
+    )
+    remove_relationship_attributes: Optional[dict[str, Any]] = Field(
+        None,
+        alias="removeRelationshipAttributes",
+        description="Map of remove relationship attributes.",
+    )
+    semantic: Optional[SaveSemantic] = Field(
+        exclude=True,
+        description=(
+            "Semantic for how this relationship should be saved, "
+            "if used in an asset request on which `.save()` is called."
+        ),
+    )

pyatlan/model/assets/asset00.py

@@ -337,9 +337,16 @@ def validate_required(self):
 
     unique_attributes: Optional[dict[str, Any]] = Field(None)
 
-    append_relationship_attributes: Optional[dict[str, Any]] = Field(None)
-    remove_relationship_attributes: Optional[dict[str, Any]] = Field(None)
-
+    append_relationship_attributes: Optional[dict[str, Any]] = Field(
+        None,
+        alias="appendRelationshipAttributes",
+        description="Map of append relationship attributes.",
+    )
+    remove_relationship_attributes: Optional[dict[str, Any]] = Field(
+        None,
+        alias="removeRelationshipAttributes",
+        description="Map of remove relationship attributes.",
+    )
     semantic: Optional[SaveSemantic] = Field(
         exclude=True,
         description=(

tests/unit/test_client.py

@@ -1340,7 +1340,10 @@ def test_process_relationship_attributes(self, glossary, term1, term2, term3):
         request_json = self.to_json(request)
         assert request_json
         assert self.SEE_ALSO in request_json["attributes"]
-        assert len(request_json["attributes"][self.SEE_ALSO]) == 2
+        replace_attributes = request_json["attributes"][self.SEE_ALSO]
+        assert len(replace_attributes) == 2
+        assert replace_attributes[0]["guid"] == term2.guid
+        assert replace_attributes[1]["guid"] == term3.guid
         assert self.APPEND not in request_json
         assert self.REMOVE not in request_json
 
@@ -1355,10 +1358,14 @@ def test_process_relationship_attributes(self, glossary, term1, term2, term3):
         request_json = self.to_json(request)
         assert request_json
         assert self.SEE_ALSO in request_json["attributes"]
-        assert len(request_json["attributes"][self.SEE_ALSO]) == 1
+        replace_attributes = request_json["attributes"][self.SEE_ALSO]
+        assert len(replace_attributes) == 1
+        assert replace_attributes[0]["guid"] == term2.guid
         assert self.APPEND in request_json
         assert self.SEE_ALSO in request_json[self.APPEND]
-        assert len(request_json[self.APPEND][self.SEE_ALSO]) == 1
+        append_attributes = request_json[self.APPEND][self.SEE_ALSO]
+        assert len(append_attributes) == 1
+        assert append_attributes[0]["guid"] == term3.guid
         assert self.REMOVE not in request_json
 
         # Test append and replace (list)
@@ -1372,10 +1379,14 @@ def test_process_relationship_attributes(self, glossary, term1, term2, term3):
         request_json = self.to_json(request)
         assert request_json
         assert self.SEE_ALSO in request_json["attributes"]
-        assert len(request_json["attributes"][self.SEE_ALSO]) == 1
+        replace_attributes = request_json["attributes"][self.SEE_ALSO]
+        assert len(replace_attributes) == 1
+        assert replace_attributes[0]["guid"] == term3.guid
         assert self.APPEND in request_json
         assert self.SEE_ALSO in request_json[self.APPEND]
-        assert len(request_json[self.APPEND][self.SEE_ALSO]) == 1
+        append_attributes = request_json[self.APPEND][self.SEE_ALSO]
+        assert len(append_attributes) == 1
+        assert append_attributes[0]["guid"] == term2.guid
         assert self.REMOVE not in request_json
 
         # Test remove and append (list)
@@ -1392,10 +1403,14 @@ def test_process_relationship_attributes(self, glossary, term1, term2, term3):
         assert request_json
         assert self.APPEND in request_json
         assert self.SEE_ALSO in request_json[self.APPEND]
-        assert len(request_json[self.APPEND][self.SEE_ALSO]) == 1
+        append_attributes = request_json[self.APPEND][self.SEE_ALSO]
+        assert len(append_attributes) == 1
+        assert append_attributes[0]["guid"] == term3.guid
         assert self.REMOVE in request_json
         assert self.SEE_ALSO in request_json[self.REMOVE]
-        assert len(request_json[self.REMOVE][self.SEE_ALSO]) == 1
+        remove_attributes = request_json[self.REMOVE][self.SEE_ALSO]
+        assert len(remove_attributes) == 1
+        assert remove_attributes[0]["guid"] == term2.guid
         assert self.SEE_ALSO not in request_json["attributes"]
 
         # Test same semantic (list)
@@ -1412,7 +1427,10 @@ def test_process_relationship_attributes(self, glossary, term1, term2, term3):
         assert request_json
         assert self.APPEND in request_json
         assert self.SEE_ALSO in request_json[self.APPEND]
-        assert len(request_json[self.APPEND][self.SEE_ALSO]) == 2
+        append_attributes = request_json[self.APPEND][self.SEE_ALSO]
+        assert len(append_attributes) == 2
+        assert append_attributes[0]["guid"] == term2.guid
+        assert append_attributes[1]["guid"] == term3.guid
         assert self.REMOVE not in request_json
         assert self.SEE_ALSO not in request_json["attributes"]
 
@@ -1422,7 +1440,8 @@ def test_process_relationship_attributes(self, glossary, term1, term2, term3):
         request_json = self.to_json(request)
         assert request_json
         assert self.SEE_ALSO in request_json["attributes"]
-        assert len(request_json["attributes"][self.SEE_ALSO]) == 0
+        replace_attributes = request_json["attributes"][self.SEE_ALSO]
+        assert len(replace_attributes) == 0
         assert self.APPEND not in request_json
         assert self.REMOVE not in request_json
 
@@ -1432,7 +1451,9 @@ def test_process_relationship_attributes(self, glossary, term1, term2, term3):
         request_json = self.to_json(request)
         assert request_json
         assert "anchor" in request_json["attributes"]
-        assert request_json["attributes"]["anchor"]
+        replace_attributes = request_json["attributes"]["anchor"]
+        assert replace_attributes
+        assert replace_attributes["guid"] == glossary.guid
         assert self.APPEND not in request_json
         assert self.REMOVE not in request_json
 
@@ -1445,7 +1466,9 @@ def test_process_relationship_attributes(self, glossary, term1, term2, term3):
         assert request_json
         assert self.APPEND in request_json
         assert "anchor" in request_json[self.APPEND]
-        assert len(request_json[self.APPEND]["anchor"]) == 1
+        append_attributes = request_json[self.APPEND]["anchor"]
+        assert len(append_attributes) == 1
+        assert append_attributes[0]["guid"] == glossary.guid
         assert self.REMOVE not in request_json
         assert "anchor" not in request_json["attributes"]
 
@@ -1458,6 +1481,8 @@ def test_process_relationship_attributes(self, glossary, term1, term2, term3):
         assert request_json
         assert self.REMOVE in request_json
         assert "anchor" in request_json[self.REMOVE]
-        assert len(request_json[self.REMOVE]["anchor"]) == 1
+        remove_attributes = request_json[self.REMOVE]["anchor"]
+        assert remove_attributes[0]["guid"] == glossary.guid
+        assert len(remove_attributes) == 1
         assert self.APPEND not in request_json
         assert "anchor" not in request_json["attributes"]