[tests] Updated unit tests as per changes

Author: Aryamanz29

pyatlan/client/atlan.py

@@ -222,7 +222,7 @@ def __init__(self, **data):
         adapter = HTTPAdapter(max_retries=self.retry)
         session.mount(HTTPS_PREFIX, adapter)
         session.mount(HTTP_PREFIX, adapter)
-        self._current_client_tls.client = None
+        AtlanClient.set_current_client(self)
 
     @property
     def cache_key(self) -> int:

tests/integration/admin_test.py

@@ -6,7 +6,6 @@
 import pytest
 from pydantic.v1 import StrictStr
 
-from pyatlan.cache.role_cache import RoleCache
 from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.group import AtlanGroup, CreateGroupResponse, GroupRequest
 from pyatlan.model.keycloak_events import AdminEventRequest, KeycloakEventRequest
@@ -35,7 +34,7 @@ def delete_group(client: AtlanClient, guid: str) -> None:
 
 
 def test_retrieve_roles(client: AtlanClient):
-    admin_role_guid = RoleCache.get_id_for_name("$admin")
+    admin_role_guid = client.role_cache.get_id_for_name("$admin")
     assert admin_role_guid
 
 

tests/integration/connection_test.py

@@ -2,7 +2,6 @@
 # Copyright 2022 Atlan Pte. Ltd.
 import pytest
 
-from pyatlan.cache.role_cache import RoleCache
 from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.assets import Connection
 from pyatlan.model.enums import AtlanConnectorType
@@ -14,7 +13,7 @@
 def create_connection(
     client: AtlanClient, name: str, connector_type: AtlanConnectorType
 ) -> Connection:
-    admin_role_guid = str(RoleCache.get_id_for_name("$admin"))
+    admin_role_guid = str(client.role_cache.get_id_for_name("$admin"))
     to_create = Connection.create(
         name=name, connector_type=connector_type, admin_roles=[admin_role_guid]
     )

tests/integration/custom_metadata_test.py

@@ -6,7 +6,6 @@
 
 import pytest
 
-from pyatlan.cache.custom_metadata_cache import CustomMetadataCache
 from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.assets import (
     Asset,
@@ -623,7 +622,7 @@ def test_search_by_any_accountable(
     term: AtlasGlossaryTerm,
 ):
     attributes = ["name", "anchor"]
-    cm_attributes = CustomMetadataCache.get_attributes_for_search_results(
+    cm_attributes = client.custom_metadata_cache.get_attributes_for_search_results(
         set_name=CM_RACI
     )
     assert cm_attributes
@@ -724,7 +723,7 @@ def test_remove_term_cm_ipr(
 def test_remove_attribute(client: AtlanClient, cm_raci: CustomMetadataDef):
     global _removal_epoch
     cm_name = CM_RACI
-    existing = CustomMetadataCache.get_custom_metadata_def(name=cm_name)
+    existing = client.custom_metadata_cache.get_custom_metadata_def(name=cm_name)
     existing_attrs = existing.attribute_defs
     updated_attrs = []
     for existing_attr in existing_attrs:
@@ -758,15 +757,15 @@ def test_remove_attribute(client: AtlanClient, cm_raci: CustomMetadataDef):
 @pytest.mark.order(after="test_remove_attribute")
 def test_retrieve_structures(client: AtlanClient, cm_raci: CustomMetadataDef):
     global _removal_epoch
-    custom_attributes = CustomMetadataCache.get_all_custom_attributes()
+    custom_attributes = client.custom_metadata_cache.get_all_custom_attributes()
     assert custom_attributes
     assert len(custom_attributes) >= 3
     assert CM_RACI in custom_attributes.keys()
     assert CM_IPR in custom_attributes.keys()
     assert CM_QUALITY in custom_attributes.keys()
     extra = _validate_raci_structure(custom_attributes.get(CM_RACI), 4)
     assert not extra
-    custom_attributes = CustomMetadataCache.get_all_custom_attributes(
+    custom_attributes = client.custom_metadata_cache.get_all_custom_attributes(
         include_deleted=True
     )
     assert custom_attributes
@@ -785,7 +784,7 @@ def test_retrieve_structures(client: AtlanClient, cm_raci: CustomMetadataDef):
 
 @pytest.mark.order(after="test_retrieve_structures")
 def test_recreate_attribute(client: AtlanClient, cm_raci: CustomMetadataDef):
-    existing = CustomMetadataCache.get_custom_metadata_def(name=CM_RACI)
+    existing = client.custom_metadata_cache.get_custom_metadata_def(name=CM_RACI)
     existing_attrs = existing.attribute_defs
     updated_attrs = []
     for existing_attr in existing_attrs:
@@ -819,7 +818,7 @@ def test_recreate_attribute(client: AtlanClient, cm_raci: CustomMetadataDef):
 def test_retrieve_structure_without_archived(
     client: AtlanClient, cm_raci: CustomMetadataDef
 ):
-    custom_attributes = CustomMetadataCache.get_all_custom_attributes()
+    custom_attributes = client.custom_metadata_cache.get_all_custom_attributes()
     assert custom_attributes
     assert len(custom_attributes) >= 3
     assert CM_RACI in custom_attributes.keys()
@@ -838,7 +837,7 @@ def test_retrieve_structure_without_archived(
 def test_retrieve_structure_with_archived(
     client: AtlanClient, cm_raci: CustomMetadataDef
 ):
-    custom_attributes = CustomMetadataCache.get_all_custom_attributes(
+    custom_attributes = client.custom_metadata_cache.get_all_custom_attributes(
         include_deleted=True
     )
     assert custom_attributes

tests/integration/test_open_lineage.py

@@ -2,7 +2,6 @@
 
 import pytest
 
-from pyatlan.cache.role_cache import RoleCache
 from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.assets import Asset, Connection, Process, SparkJob
 from pyatlan.model.audit import AuditSearchRequest
@@ -17,7 +16,7 @@
 
 @pytest.fixture(scope="module")
 def connection(client: AtlanClient):
-    admin_role_guid = RoleCache.get_id_for_name("$admin")
+    admin_role_guid = client.role_cache.get_id_for_name("$admin")
     assert admin_role_guid
     response = client.open_lineage.create_connection(
         name=MODULE_NAME, admin_roles=[admin_role_guid]

tests/integration/test_sql_assets.py

@@ -5,7 +5,6 @@
 
 import pytest
 
-from pyatlan.cache.role_cache import RoleCache
 from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.assets import (
     Asset,
@@ -76,7 +75,7 @@ def test_create(
         client: AtlanClient,
         upsert: Callable[[Asset], AssetMutationResponse],
     ):
-        role = RoleCache.get_id_for_name("$admin")
+        role = client.role_cache.get_id_for_name("$admin")
         assert role
         connection_name = TestId.make_unique("INT")
         c = Connection.create(

tests/unit/model/badge_test.py

@@ -1,5 +1,8 @@
+from unittest.mock import MagicMock
+
 import pytest
 
+from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.assets import Badge
 from pyatlan.model.enums import BadgeComparisonOperator, BadgeConditionColor
 from pyatlan.model.structs import BadgeCondition
@@ -23,6 +26,33 @@
 BADGE_METADATA_ATTRIBUTE = f"{CM_ID}.{CM_ATTR_ID}"
 
 
+@pytest.fixture(autouse=True)
+def set_env(monkeypatch):
+    monkeypatch.setenv("ATLAN_BASE_URL", "https://test.atlan.com")
+    monkeypatch.setenv("ATLAN_API_KEY", "test-api-key")
+
+
+@pytest.fixture()
+def client():
+    return AtlanClient()
+
+
+@pytest.fixture()
+def current_client(client, monkeypatch):
+    monkeypatch.setattr(
+        AtlanClient,
+        "get_current_client",
+        lambda: client,
+    )
+
+
+@pytest.fixture()
+def mock_cm_cache(current_client, monkeypatch):
+    mock_cache = MagicMock()
+    monkeypatch.setattr(AtlanClient, "custom_metadata_cache", mock_cache)
+    return mock_cache
+
+
 @pytest.mark.parametrize(
     "name, cm_name, cm_attribute, badge_conditions, message",
     [
@@ -47,9 +77,9 @@ def test_create_when_required_parameters_are_missing_raises_value_error(
         )
 
 
-def test_create(mock_custom_metadata_cache):
-    mock_custom_metadata_cache.get_attr_id_for_name.return_value = CM_ATTR_ID
-    mock_custom_metadata_cache.get_id_for_name.return_value = CM_ID
+def test_create(mock_cm_cache):
+    mock_cm_cache.get_attr_id_for_name.return_value = CM_ATTR_ID
+    mock_cm_cache.get_id_for_name.return_value = CM_ID
 
     badge = Badge.create(
         name=BADGE_NAME,

tests/unit/test_atlan_tag_name.py

@@ -4,6 +4,7 @@
 from pydantic.v1 import parse_obj_as
 
 import pyatlan.cache.atlan_tag_cache
+from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.assets import Purpose
 from pyatlan.model.constants import DELETED_
 from pyatlan.model.core import AtlanTagName
@@ -13,8 +14,30 @@
 GOOD_ATLAN_TAG_NAME = "PII"
 
 
-def test_init_with_bad_atlan_tag_name_raises_value_error(monkeypatch):
-    def get_id_for_name(_):
+@pytest.fixture(autouse=True)
+def set_env(monkeypatch):
+    monkeypatch.setenv("ATLAN_BASE_URL", "https://test.atlan.com")
+    monkeypatch.setenv("ATLAN_API_KEY", "test-api-key")
+
+
+@pytest.fixture()
+def client():
+    return AtlanClient()
+
+
+@pytest.fixture()
+def current_client(client, monkeypatch):
+    monkeypatch.setattr(
+        AtlanClient,
+        "get_current_client",
+        lambda: client,
+    )
+
+
+def test_init_with_bad_atlan_tag_name_raises_value_error(
+    current_client: AtlanClient, monkeypatch
+):
+    def get_id_for_name(_, __):
         return None
 
     monkeypatch.setattr(
@@ -29,8 +52,8 @@ def get_id_for_name(_):
 
 
 @pytest.fixture()
-def good_atlan_tag(monkeypatch):
-    def get_id_for_name(value):
+def good_atlan_tag(current_client: AtlanClient, monkeypatch):
+    def get_id_for_name(_, value):
         return ATLAN_TAG_ID
 
     monkeypatch.setattr(
@@ -41,8 +64,8 @@ def get_id_for_name(value):
     return AtlanTagName(GOOD_ATLAN_TAG_NAME)
 
 
-def test_init_with_good_name(monkeypatch):
-    def get_id_for_name(value):
+def test_init_with_good_name(current_client: AtlanClient, monkeypatch):
+    def get_id_for_name(_, value):
         assert value == GOOD_ATLAN_TAG_NAME
         return GOOD_ATLAN_TAG_NAME
 
@@ -60,13 +83,15 @@ def get_id_for_name(value):
 
 
 def test_convert_to_display_text_when_atlan_tag_passed_returns_same_atlan_tag(
-    monkeypatch, good_atlan_tag
+    good_atlan_tag,
 ):
     assert good_atlan_tag is AtlanTagName._convert_to_display_text(good_atlan_tag)
 
 
-def test_convert_to_display_text_when_bad_string(monkeypatch):
-    def get_name_for_id(_):
+def test_convert_to_display_text_when_bad_string(
+    current_client: AtlanClient, monkeypatch
+):
+    def get_name_for_id(_, __):
         return None
 
     monkeypatch.setattr(
@@ -81,11 +106,11 @@ def get_name_for_id(_):
     )
 
 
-def test_convert_to_display_text_when_id(monkeypatch):
-    def get_name_for_id(value):
+def test_convert_to_display_text_when_id(current_client: AtlanClient, monkeypatch):
+    def get_name_for_id(_, __):
         return GOOD_ATLAN_TAG_NAME
 
-    def get_id_for_name(value):
+    def get_id_for_name(_, value):
         assert value == GOOD_ATLAN_TAG_NAME
         return GOOD_ATLAN_TAG_NAME
 
@@ -105,15 +130,15 @@ def get_id_for_name(value):
     assert str(sut) == GOOD_ATLAN_TAG_NAME
 
 
-def test_json_encode_atlan_tag(monkeypatch, good_atlan_tag):
+def test_json_encode_atlan_tag(good_atlan_tag):
     assert AtlanTagName.json_encode_atlan_tag(good_atlan_tag) == ATLAN_TAG_ID
 
 
-def test_asset_tag_name_field_deserialization(monkeypatch):
-    def get_name_for_id(_):
+def test_asset_tag_name_field_deserialization(current_client: AtlanClient, monkeypatch):
+    def get_name_for_id(_, __):
         return None
 
-    def get_id_for_name(_):
+    def get_id_for_name(_, __):
         return None
 
     monkeypatch.setattr(