[tests] Use token_client (retry=0) to avoid extra token generation

Author: Aryamanz29

tests/integration/aio/conftest.py

@@ -14,6 +14,7 @@
 from pyatlan.model.atlan_image import AtlanImage
 from pyatlan.model.enums import AtlanConnectorType, AtlanTagColor, CertificateStatus
 from pyatlan.model.typedef import AtlanTagDef
+from pyatlan.client.atlan import DEFAULT_RETRY
 from tests.integration.aio.test_connection import create_connection_async
 from tests.integration.aio.utils import delete_asset_async
 from tests.integration.client import TestId
@@ -36,6 +37,13 @@ async def client():
     client = AsyncAtlanClient()
     yield client
 
+@pytest_asyncio.fixture(scope="module")
+async def token_client():
+    """Async Atlan client fixture for api token integration tests."""
+    DEFAULT_RETRY.total = 0
+    client = AsyncAtlanClient(retry=DEFAULT_RETRY)
+    yield client
+
 
 @pytest_asyncio.fixture(scope="module")
 async def connection(client: AsyncAtlanClient) -> AsyncGenerator[Connection, None]:

tests/integration/aio/test_client.py

@@ -77,34 +77,34 @@ async def current_user(client: AsyncAtlanClient) -> UserMinimalResponse:
 
 
 @pytest_asyncio.fixture(scope="module")
-async def token(client: AsyncAtlanClient) -> AsyncGenerator[ApiToken, None]:
+async def token(token_client: AsyncAtlanClient) -> AsyncGenerator[ApiToken, None]:
     token = None
     try:
-        token = await create_token_async(client, MODULE_NAME)
+        token = await create_token_async(token_client, MODULE_NAME)
         yield token
     finally:
-        await delete_token_async(client, token)
+        await delete_token_async(token_client, token)
 
 
 @pytest_asyncio.fixture(scope="module")
-async def expired_token(client: AsyncAtlanClient) -> AsyncGenerator[ApiToken, None]:
+async def expired_token(token_client: AsyncAtlanClient) -> AsyncGenerator[ApiToken, None]:
     token = None
     try:
-        token = await client.token.create(f"{MODULE_NAME}-expired", validity_seconds=1)
+        token = await token_client.token.create(f"{MODULE_NAME}-expired", validity_seconds=1)
         time.sleep(5)
         yield token
     finally:
-        await delete_token_async(client, token)
+        await delete_token_async(token_client, token)
 
 
 @pytest_asyncio.fixture(scope="module")
-async def argo_fake_token(client: AsyncAtlanClient) -> AsyncGenerator[ApiToken, None]:
+async def argo_fake_token(token_client: AsyncAtlanClient) -> AsyncGenerator[ApiToken, None]:
     token = None
     try:
-        token = await client.token.create(f"{MODULE_NAME}-fake-argo")
+        token = await token_client.token.create(f"{MODULE_NAME}-fake-argo")
         yield token
     finally:
-        await delete_token_async(client, token)
+        await delete_token_async(token_client, token)
 
 
 @pytest_asyncio.fixture(scope="module")

tests/integration/aio/utils.py

@@ -19,22 +19,22 @@
 LOGGER = logging.getLogger(__name__)
 
 
-async def create_token_async(client: AsyncAtlanClient, name: str) -> ApiToken:
+async def create_token_async(token_client: AsyncAtlanClient, name: str) -> ApiToken:
     """Create an API token asynchronously."""
-    token = await client.token.create(name)
+    token = await token_client.token.create(name)
     return token
 
 
 async def delete_token_async(
-    client: AsyncAtlanClient, token: Optional[ApiToken] = None
+    token_client: AsyncAtlanClient, token: Optional[ApiToken] = None
 ):
     """Delete an API token asynchronously."""
     # If there is a partial failure on the server side
     # and the token is still visible in the Atlan UI,
     # in that case, the create method may not return a token.
     # We should retrieve the list of all tokens and delete them here.
     if not token:
-        tokens_response = await client.token.get()
+        tokens_response = await token_client.token.get()
         tokens = tokens_response.records
         assert tokens
         delete_tokens = [
@@ -44,11 +44,11 @@ async def delete_token_async(
         ]
         for token in delete_tokens:
             assert token and token.guid
-            await client.token.purge(token.guid)
+            await token_client.token.purge(token.guid)
         return
     # In case of no partial failure, directly delete the token
     if token.guid:
-        await client.token.purge(token.guid)
+        await token_client.token.purge(token.guid)
 
 
 async def delete_asset_async(

tests/integration/client.py

@@ -8,6 +8,7 @@
 from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.enums import AtlanDeleteType
 from pyatlan.model.response import A
+from pyatlan.client.atlan import DEFAULT_RETRY
 
 LOGGER = logging.getLogger(__name__)
 
@@ -31,6 +32,12 @@ def client() -> Generator[AtlanClient, None, None]:
 
     yield client
 
+@pytest.fixture(scope="module")
+def token_client() -> Generator[AtlanClient, None, None]:
+    DEFAULT_RETRY.total = 0
+    client = AtlanClient(retry=DEFAULT_RETRY)
+    yield client
+
 
 def delete_asset(
     client: AtlanClient,

tests/integration/purpose_test.py

@@ -48,18 +48,18 @@ def snowflake_column_qn(snowflake_conn):
 
 
 @pytest.fixture(scope="module")
-def token(client: AtlanClient) -> Generator[ApiToken, None, None]:
+def token(token_client: AtlanClient) -> Generator[ApiToken, None, None]:
     token = None
     try:
-        token = client.token.create(API_TOKEN_NAME)
+        token = token_client.token.create(API_TOKEN_NAME)
         assert token
         assert token.guid
         assert token.display_name
         # After creating the token, assign it to the
         # "Data Assets" persona to grant it query access
-        persona = client.asset.find_personas_by_name(PERSONA_NAME)[0]
+        persona = token_client.asset.find_personas_by_name(PERSONA_NAME)[0]
         assert persona.qualified_name
-        client.token.update(
+        token_client.token.update(
             guid=token.guid,
             display_name=token.display_name,
             personas={persona.qualified_name},
@@ -68,7 +68,7 @@ def token(client: AtlanClient) -> Generator[ApiToken, None, None]:
         # its associated personas -- will leave that to later...
         yield token
     finally:
-        delete_token(client, token)
+        delete_token(token_client, token)
 
 
 @pytest.fixture(scope="module")

tests/integration/requests_test.py

@@ -10,18 +10,18 @@
 API_TOKEN_NAME = f"{MODULE_NAME}"
 
 
-def create_token(client: AtlanClient, name: str) -> ApiToken:
-    t = client.token.create(name)
+def create_token(token_client: AtlanClient, name: str) -> ApiToken:
+    t = token_client.token.create(name)
     return t
 
 
-def delete_token(client: AtlanClient, token: Optional[ApiToken] = None):
+def delete_token(token_client: AtlanClient, token: Optional[ApiToken] = None):
     # If there is a partial failure on the server side
     # and the token is still visible in the Atlan UI,
     # in that case, the create method may not return a token.
     # We should retrieve the list of all tokens and delete them here.
     if not token:
-        tokens = client.token.get().records
+        tokens = token_client.token.get().records
         assert tokens
         delete_tokens = [
             token
@@ -30,20 +30,20 @@ def delete_token(client: AtlanClient, token: Optional[ApiToken] = None):
         ]
         for token in delete_tokens:
             assert token and token.guid
-            client.token.purge(token.guid)
+            token_client.token.purge(token.guid)
         return
     # In case of no partial failure, directly delete the token
-    token.guid and client.token.purge(token.guid)
+    token.guid and token_client.token.purge(token.guid)
 
 
 @pytest.fixture(scope="module")
-def token(client: AtlanClient) -> Generator[ApiToken, None, None]:
+def token(token_client: AtlanClient) -> Generator[ApiToken, None, None]:
     token = None
     try:
-        token = create_token(client, API_TOKEN_NAME)
+        token = create_token(token_client, API_TOKEN_NAME)
         yield token
     finally:
-        delete_token(client, token)
+        delete_token(token_client, token)
 
 
 def test_create_token(client: AtlanClient, token: ApiToken):

tests/integration/test_client.py

@@ -66,34 +66,34 @@
 
 
 @pytest.fixture(scope="module")
-def token(client: AtlanClient) -> Generator[ApiToken, None, None]:
+def token(token_client: AtlanClient) -> Generator[ApiToken, None, None]:
     token = None
     try:
-        token = create_token(client, MODULE_NAME)
+        token = create_token(token_client, MODULE_NAME)
         yield token
     finally:
-        delete_token(client, token)
+        delete_token(token_client, token)
 
 
 @pytest.fixture(scope="module")
-def expired_token(client: AtlanClient) -> Generator[ApiToken, None, None]:
+def expired_token(token_client: AtlanClient) -> Generator[ApiToken, None, None]:
     token = None
     try:
-        token = client.token.create(f"{MODULE_NAME}-expired", validity_seconds=1)
+        token = token_client.token.create(f"{MODULE_NAME}-expired", validity_seconds=1)
         time.sleep(5)
         yield token
     finally:
-        delete_token(client, token)
+        delete_token(token_client, token)
 
 
 @pytest.fixture(scope="module")
-def argo_fake_token(client: AtlanClient) -> Generator[ApiToken, None, None]:
+def argo_fake_token(token_client: AtlanClient) -> Generator[ApiToken, None, None]:
     token = None
     try:
-        token = client.token.create(f"{MODULE_NAME}-fake-argo")
+        token = token_client.token.create(f"{MODULE_NAME}-fake-argo")
         yield token
     finally:
-        delete_token(client, token)
+        delete_token(token_client, token)
 
 
 @pytest.fixture(scope="module")