Merge branch 'main' into FT-896

Author: Aryamanz29

pyatlan/pkg/utils.py

@@ -82,7 +82,6 @@ def get_client(impersonate_user_id: str) -> AtlanClient:
     elif user_id:
         LOGGER.info("No API token found, attempting to impersonate user: %s", user_id)
         client = AtlanClient(base_url=base_url, api_key="")
-        client._user_id = user_id
         api_key = client.impersonate.user(user_id=user_id)
     else:
         LOGGER.info(
@@ -91,7 +90,10 @@ def get_client(impersonate_user_id: str) -> AtlanClient:
         client = AtlanClient(base_url=base_url, api_key="")
         api_key = client.impersonate.escalate()
 
-    return AtlanClient(base_url=base_url, api_key=api_key)
+    client = AtlanClient(base_url=base_url, api_key=api_key)
+    if user_id:
+        client._user_id = user_id
+    return client
 
 
 def set_package_ops(run_time_config: RuntimeConfig) -> AtlanClient:

tests/integration/custom_package_test.py

@@ -6,9 +6,10 @@
 import pytest
 
 from pyatlan.client.atlan import AtlanClient
+from pyatlan.client.impersonate import ImpersonationClient
 from pyatlan.pkg.models import CustomPackage, generate
 from pyatlan.pkg.ui import UIConfig, UIStep
-from pyatlan.pkg.utils import set_package_headers
+from pyatlan.pkg.utils import get_client, set_package_headers
 from pyatlan.pkg.widgets import TextInput
 
 
@@ -21,6 +22,9 @@ def mock_pkg_env():
             "X_ATLAN_AGENT_ID": "agent_id_value",
             "X_ATLAN_AGENT_PACKAGE_NAME": "package_name_value",
             "X_ATLAN_AGENT_WORKFLOW_ID": "workflow_id_value",
+            "X_ATLAN_AGENT_WORKFLOW_ID": "workflow_id_value",
+            "CLIENT_ID": "client_id_value",
+            "CLIENT_SECRET": "client_secret_value",
         },
         clear=True,
     ):
@@ -104,3 +108,16 @@ def test_set_package_headers(client: AtlanClient, mock_pkg_env):
     }
     mock_client.update_headers.assert_called_once_with(expected_headers)
     assert updated_client == mock_client
+
+
+@patch.object(ImpersonationClient, "user", return_value="some-api-key")
+def test_get_client_user_id_handling(
+    mock_impersonate_client,
+    mock_pkg_env,
+    client: AtlanClient,
+):
+    updated_client = get_client(impersonate_user_id="test-user-id")
+    assert updated_client
+    assert updated_client.base_url
+    assert updated_client.api_key == "some-api-key"
+    assert updated_client._user_id == "test-user-id"

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"