Add replace_terms method to AtlanClient

Author: ErnestoLoma

pyatlan/client/atlan.py

@@ -634,7 +634,33 @@ def append_terms(
                 term for term in existing_terms if term.relationship_status != "DELETED"
             )
         replacement_terms.extend(terms)
-        asset.attributes.meanings = replacement_terms
+        asset.terms = replacement_terms
+        response = self.upsert(entity=asset)
+        if assets := response.assets_updated(asset_type=asset_type):
+            return assets[0]
+        return asset
+
+    @validate_arguments()
+    def replace_terms(
+        self,
+        asset_type: Type[A],
+        terms: list[AtlasGlossaryTerm],
+        guid: Optional[str] = None,
+        qualified_name: Optional[str] = None,
+    ) -> A:
+        if guid:
+            if qualified_name:
+                raise ValueError(
+                    "Either guid or qualified_name can be be specified not both"
+                )
+            asset = self.get_asset_by_guid(guid=guid, asset_type=asset_type)
+        elif qualified_name:
+            asset = self.get_asset_by_qualified_name(
+                qualified_name=qualified_name, asset_type=asset_type
+            )
+        else:
+            raise ValueError("Either guid or qualified name must be specified")
+        asset.terms = terms
         response = self.upsert(entity=asset)
         if assets := response.assets_updated(asset_type=asset_type):
             return assets[0]

tests/integration/test_client.py

@@ -126,3 +126,61 @@ def test_append_terms_using_ref_by_guid_for_term(
     database = client.get_asset_by_guid(guid=database.guid, asset_type=Database)
     assert len(database.terms) == 1
     assert database.terms[0].guid == term.guid
+
+
+def test_replace_a_term(
+    client: AtlanClient,
+    make_term: Callable[[str], AtlasGlossaryTerm],
+    database: Database,
+):
+    original_term = make_term("Term1")
+    assert (
+        database := client.append_terms(
+            qualified_name=database.qualified_name,
+            asset_type=Database,
+            terms=[AtlasGlossaryTerm.ref_by_guid(guid=original_term.guid)],
+        )
+    )
+
+    replacemant_term = make_term("Term2")
+    assert (
+        database := client.replace_terms(
+            guid=database.guid, asset_type=Database, terms=[replacemant_term]
+        )
+    )
+
+    database = client.get_asset_by_guid(guid=database.guid, asset_type=Database)
+    assert len(database.terms) == 2
+    deleted_terms = [t for t in database.terms if t.relationship_status == "DELETED"]
+    assert len(deleted_terms) == 1
+    assert deleted_terms[0].guid == original_term.guid
+    active_terms = [t for t in database.terms if t.relationship_status != "DELETED"]
+    assert len(active_terms) == 1
+    assert active_terms[0].guid == replacemant_term.guid
+
+
+def test_replace_all_term(
+    client: AtlanClient,
+    make_term: Callable[[str], AtlasGlossaryTerm],
+    database: Database,
+):
+    original_term = make_term("Term1")
+    assert (
+        database := client.append_terms(
+            qualified_name=database.qualified_name,
+            asset_type=Database,
+            terms=[AtlasGlossaryTerm.ref_by_guid(guid=original_term.guid)],
+        )
+    )
+
+    assert (
+        database := client.replace_terms(
+            guid=database.guid, asset_type=Database, terms=[]
+        )
+    )
+
+    database = client.get_asset_by_guid(guid=database.guid, asset_type=Database)
+    assert len(database.terms) == 1
+    deleted_terms = [t for t in database.terms if t.relationship_status == "DELETED"]
+    assert len(deleted_terms) == 1
+    assert deleted_terms[0].guid == original_term.guid

tests/unit/test_client.py

@@ -152,3 +152,76 @@ def test_append_with_valid_guid_when_terms_present_returns_asset_with_combined_t
         assert len(updated_terms) == 2
         assert exisiting_term in updated_terms
         assert new_term in updated_terms
+
+
+@pytest.mark.parametrize(
+    "guid, qualified_name, asset_type, terms, message",
+    [
+        (
+            None,
+            None,
+            Table,
+            [AtlasGlossaryTerm()],
+            "Either guid or qualified name must be specified",
+        ),
+        (
+            "123",
+            None,
+            None,
+            [AtlasGlossaryTerm()],
+            "1 validation error for ReplaceTerms\\nasset_type\\n  none is not an allowed value ",
+        ),
+        (
+            "123",
+            "default/abc",
+            Table,
+            [AtlasGlossaryTerm()],
+            "Either guid or qualified_name can be be specified not both",
+        ),
+        (
+            "123",
+            None,
+            Table,
+            None,
+            "1 validation error for ReplaceTerms\\nterms\\n  none is not an allowed value ",
+        ),
+    ],
+)
+def test_replace_terms_with_invalid_parameter_raises_valueerror(
+    guid,
+    qualified_name,
+    asset_type,
+    terms,
+    message,
+    monkeypatch,
+):
+    monkeypatch.setenv("ATLAN_BASE_URL", "https://name.atlan.com")
+    monkeypatch.setenv("ATLAN_API_KEY", "abkj")
+    client = AtlanClient()
+
+    with pytest.raises(ValueError, match=message):
+        client.replace_terms(
+            guid=guid, qualified_name=qualified_name, asset_type=asset_type, terms=terms
+        )
+
+
+def test_replace_terms(
+    monkeypatch,
+):
+    monkeypatch.setenv("ATLAN_BASE_URL", "https://name.atlan.com")
+    monkeypatch.setenv("ATLAN_API_KEY", "abkj")
+    asset_type = Table
+    with patch.multiple(
+        AtlanClient, get_asset_by_guid=DEFAULT, upsert=DEFAULT
+    ) as mock_methods:
+        table = Table()
+        mock_methods["get_asset_by_guid"].return_value = table
+        mock_methods["upsert"].return_value.assets_updated.return_value = [table]
+        client = AtlanClient()
+        guid = "123"
+        terms = [AtlasGlossaryTerm()]
+
+        assert (
+            asset := client.replace_terms(guid=guid, asset_type=asset_type, terms=terms)
+        )
+        assert asset.terms == terms