FT-910: Fixed _user_id handling in get_client() method

Author: Aryamanz29

Diff for: 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:

Diff for: 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"