Merge pull request #476 from atlanhq/FT-897

FT-897: Fixing the unit tests for append terms, replace terms and remove terms

Author: Aryamanz29

tests/unit/test_client.py

@@ -336,6 +336,97 @@ def test_append_terms_asset_retrieval_errors(
         )
 
 
+def test_append_with_valid_guid_and_no_terms_returns_asset():
+    asset_type = Table
+    table = Table()
+    table.name = "table-test"
+    table.qualified_name = "table_qn"
+
+    terms = []
+
+    with patch(
+        "pyatlan.model.fluent_search.FluentSearch.execute"
+    ) as mock_execute, patch("pyatlan.client.asset.AssetClient.save") as mock_save:
+        mock_execute.return_value.current_page = lambda: [table]
+
+        mock_save.return_value.assets_updated.return_value = [table]
+
+        client = AtlanClient()
+        guid = "123"
+
+        asset = client.asset.append_terms(guid=guid, asset_type=asset_type, terms=terms)
+
+        assert asset == table
+        assert asset.assigned_terms is None
+        mock_execute.assert_called_once()
+        mock_save.assert_called_once()
+
+
+def test_append_with_valid_guid_when_no_terms_present_returns_asset_with_given_terms():
+    asset_type = Table
+    table = Table()
+    table.name = "table-test"
+    table.qualified_name = "table_qn"
+
+    terms = [AtlasGlossaryTerm(qualified_name="term1")]
+
+    with patch(
+        "pyatlan.model.fluent_search.FluentSearch.execute"
+    ) as mock_execute, patch("pyatlan.client.asset.AssetClient.save") as mock_save:
+        mock_execute.return_value.current_page = lambda: [table]
+
+        def mock_save_side_effect(entity):
+            entity.assigned_terms = terms
+            return Mock(assets_updated=lambda asset_type: [entity])
+
+        mock_save.side_effect = mock_save_side_effect
+
+        client = AtlanClient()
+        guid = "123"
+        asset = client.asset.append_terms(guid=guid, asset_type=asset_type, terms=terms)
+
+        assert asset.assigned_terms == terms
+        mock_execute.assert_called_once()
+        mock_save.assert_called_once()
+
+
+def test_append_with_valid_guid_when_terms_present_returns_asset_with_combined_terms():
+    asset_type = Table
+    table = Table()
+    table.name = "table-test"
+    table.qualified_name = "table_qn"
+
+    exisiting_term = AtlasGlossaryTerm()
+    table.attributes.meanings = [exisiting_term]
+
+    new_term = AtlasGlossaryTerm(qualified_name="new_term")
+    terms = [new_term]
+
+    with patch(
+        "pyatlan.model.fluent_search.FluentSearch.execute"
+    ) as mock_execute, patch("pyatlan.client.asset.AssetClient.save") as mock_save:
+        mock_execute.return_value.current_page = lambda: [table]
+
+        def mock_save_side_effect(entity):
+            entity.assigned_terms = table.attributes.meanings + terms
+            return Mock(assets_updated=lambda asset_type: [entity])
+
+        mock_save.side_effect = mock_save_side_effect
+
+        client = AtlanClient()
+        guid = "123"
+
+        asset = client.asset.append_terms(guid=guid, asset_type=asset_type, terms=terms)
+
+        updated_terms = asset.assigned_terms
+        assert updated_terms is not None
+        assert len(updated_terms) == 2
+        assert exisiting_term in updated_terms
+        assert new_term in updated_terms
+        mock_execute.assert_called_once()
+        mock_save.assert_called_once()
+
+
 @pytest.mark.parametrize(
     "guid, qualified_name, asset_type, assigned_terms, expected_message, expected_error",
     [
@@ -437,6 +528,40 @@ def test_replace_terms_asset_retrieval_errors(
         )
 
 
+def test_replace_terms():
+    asset_type = Table
+    table = Table()
+    table.name = "table-test"
+    table.qualified_name = "table_qn"
+
+    exisiting_term = AtlasGlossaryTerm()
+    table.attributes.meanings = [exisiting_term]
+
+    terms = [AtlasGlossaryTerm(qualified_name="new_term")]
+
+    with patch(
+        "pyatlan.model.fluent_search.FluentSearch.execute"
+    ) as mock_execute, patch("pyatlan.client.asset.AssetClient.save") as mock_save:
+        mock_execute.return_value.current_page = lambda: [table]
+
+        def mock_save_side_effect(entity):
+            entity.assigned_terms = terms
+            return Mock(assets_updated=lambda asset_type: [entity])
+
+        mock_save.side_effect = mock_save_side_effect
+
+        client = AtlanClient()
+        guid = "123"
+
+        asset = client.asset.replace_terms(
+            guid=guid, asset_type=asset_type, terms=terms
+        )
+
+        assert asset.assigned_terms == terms
+        mock_execute.assert_called_once()
+        mock_save.assert_called_once()
+
+
 @pytest.mark.parametrize(
     "guid, qualified_name, asset_type, assigned_terms, expected_message, expected_error",
     [
@@ -538,6 +663,48 @@ def test_remove_terms_asset_retrieval_errors(
         )
 
 
+def test_remove_with_valid_guid_when_terms_present_returns_asset_with_terms_removed():
+    asset_type = Table
+    table = Table()
+    table.name = "table-test"
+    table.qualified_name = "table_qn"
+
+    existing_term = AtlasGlossaryTerm(
+        qualified_name="term_to_remove", guid="b4113341-251b-4adc-81fb-2420501c30e6"
+    )
+    other_term = AtlasGlossaryTerm(
+        qualified_name="other_term", guid="b267858d-8316-4c41-a56a-6e9b840cef4a"
+    )
+    table.attributes.meanings = [existing_term, other_term]
+
+    with patch(
+        "pyatlan.model.fluent_search.FluentSearch.execute"
+    ) as mock_execute, patch("pyatlan.client.asset.AssetClient.save") as mock_save:
+        mock_execute.return_value.current_page = lambda: [table]
+
+        def mock_save_side_effect(entity):
+            entity.assigned_terms = [
+                t for t in table.attributes.meanings if t != existing_term
+            ]
+            return Mock(assets_updated=lambda asset_type: [entity])
+
+        mock_save.side_effect = mock_save_side_effect
+
+        client = AtlanClient()
+        guid = "123"
+
+        asset = client.asset.remove_terms(
+            guid=guid, asset_type=asset_type, terms=[existing_term]
+        )
+
+        updated_terms = asset.assigned_terms
+        assert updated_terms is not None
+        assert len(updated_terms) == 1
+        assert other_term in updated_terms
+        mock_execute.assert_called_once()
+        mock_save.assert_called_once()
+
+
 def test_register_client_with_bad_parameter_raises_value_error(client):
     with pytest.raises(
         InvalidRequestError, match="client must be an instance of AtlanClient"