test: harden admin, audit, and AI asset integration tests against tenant state and index lag

tests/integration/admin_test.py

@@ -21,7 +21,7 @@
 
 EMAIL_DOMAIN = "@atlan.com"
 
-_default_group_count: int = 0
+_user_group_count_baseline: int = 0
 
 
 def create_group(client: AtlanClient, name: str) -> CreateGroupResponse:
@@ -41,9 +41,11 @@ def test_retrieve_roles(client: AtlanClient):
 
 @pytest.fixture(scope="module")
 def group(client: AtlanClient) -> Generator[CreateGroupResponse, None, None]:
+    global _user_group_count_baseline
     to_create = AtlanGroup.create(GROUP_NAME)
     fixed_user = client.user.get_by_username(FIXED_USER)
     assert fixed_user
+    _user_group_count_baseline = fixed_user.group_count or 0
     g = client.group.create(group=to_create, user_ids=[str(fixed_user.id)])
     yield g
     delete_group(client, g.group)
@@ -84,13 +86,9 @@ def test_create_group(client: AtlanClient, group: CreateGroupResponse):
 
 
 def test_retrieve_all_groups(client: AtlanClient, group: CreateGroupResponse):
-    global _default_group_count
     groups = client.group.get_all()
     assert groups.records
     assert len(groups.records) >= 1
-    for group1 in groups.records:
-        if group1.is_default():
-            _default_group_count += 1
 
 
 def test_group_get_all_pagination(client: AtlanClient):
@@ -238,14 +236,13 @@ def test_user_get_by_email_and_emails_pagination(client: AtlanClient):
 
 @pytest.mark.order(after="test_retrieve_all_groups")
 def test_retrieve_existing_user(client: AtlanClient, group: CreateGroupResponse):
-    global _default_group_count
     all_users = client.user.get_all()
     assert all_users.records
     assert len(all_users.records) >= 1  # type: ignore
     user1 = client.user.get_by_username(FIXED_USER)
     assert user1
     assert user1.id
-    assert user1.group_count == 1 + _default_group_count
+    assert user1.group_count == _user_group_count_baseline + 1
     response = client.user.get_by_usernames(usernames=[FIXED_USER])
     assert response
     assert response.records is not None
@@ -332,16 +329,12 @@ def test_final_user_state(
     client: AtlanClient,
     group: CreateGroupResponse,
 ):
-    global _default_group_count
     fixed_user = client.user.get_by_username(FIXED_USER)
     assert fixed_user
     assert fixed_user.id
     response = client.user.get_groups(fixed_user.id)
-    assert (
-        response.records is None
-        or len(response.records) == 0
-        or len(response.records) == _default_group_count
-    )
+    group_ids = {g.id for g in response.records or []}
+    assert group.group not in group_ids
 
 
 @pytest.mark.order(after="test_final_user_state")

tests/integration/ai_asset_test.py

@@ -3,9 +3,16 @@
 from typing import Generator
 
 import pytest
+from tenacity import (
+    RetryError,
+    retry,
+    retry_if_result,
+    stop_after_attempt,
+    wait_exponential,
+)
 
 from pyatlan.client.atlan import AtlanClient
-from pyatlan.model.assets import AIApplication, AIModel, Asset, Connection
+from pyatlan.model.assets import AIApplication, AIModel, Asset, Connection, Process
 from pyatlan.model.enums import (
     AIApplicationDevelopmentStage,
     AIDatasetType,
@@ -112,45 +119,45 @@ def test_update_ai_assets(
     _update_ai_model(client, ai_model)
 
 
+def _process_with_io(client: AtlanClient, created: Process) -> Process:
+    """
+    The mutation response does not reliably echo the process inputs/outputs
+    relationships, so re-fetch the process until they are materialized.
+    """
+    if created.inputs and created.outputs:
+        return created
+
+    @retry(
+        retry=retry_if_result(lambda process: not (process.inputs and process.outputs)),
+        stop=stop_after_attempt(10),
+        wait=wait_exponential(multiplier=1, min=2, max=10),
+    )
+    def _fetch() -> Process:
+        return client.asset.get_by_guid(
+            created.guid, asset_type=Process, ignore_relationships=False
+        )
+
+    try:
+        return _fetch()
+    except RetryError as err:
+        return err.last_attempt.result()
+
+
 def _assert_response_processes_creator(
-    mutation_response, asset_list, ai_dataset_type, process_sum, ai_model
+    client, mutation_response, asset_list, ai_dataset_type, process_sum, ai_model
 ):
     for i in range(len(asset_list)):
-        assert mutation_response.mutated_entities.CREATE[i + process_sum]
-        assert (
-            mutation_response.mutated_entities.CREATE[i + process_sum].ai_dataset_type  # type: ignore
-            == ai_dataset_type
-        )
+        created = mutation_response.mutated_entities.CREATE[i + process_sum]
+        assert created
+        assert created.ai_dataset_type == ai_dataset_type
+        process = _process_with_io(client, created)
+        assert process.inputs and process.outputs
         if ai_dataset_type == AIDatasetType.OUTPUT:
-            assert (
-                mutation_response.mutated_entities.CREATE[i + process_sum].inputs  # type: ignore
-                and mutation_response.mutated_entities.CREATE[i + process_sum]
-                .inputs[0]
-                .guid
-                == ai_model.guid  # type: ignore
-            )
-            assert (
-                mutation_response.mutated_entities.CREATE[i + process_sum].outputs  # type: ignore
-                and mutation_response.mutated_entities.CREATE[i + process_sum]
-                .outputs[0]
-                .guid  # type: ignore
-                == asset_list[i].guid
-            )
+            assert process.inputs[0].guid == ai_model.guid
+            assert process.outputs[0].guid == asset_list[i].guid
         else:
-            assert (
-                mutation_response.mutated_entities.CREATE[i + process_sum].inputs  # type: ignore
-                and mutation_response.mutated_entities.CREATE[i + process_sum]
-                .inputs[0]
-                .guid
-                == asset_list[i].guid  # type: ignore
-            )
-            assert (
-                mutation_response.mutated_entities.CREATE[i + process_sum].outputs  # type: ignore
-                and mutation_response.mutated_entities.CREATE[i + process_sum]
-                .outputs[0]
-                .guid  # type: ignore
-                == ai_model.guid
-            )
+            assert process.inputs[0].guid == asset_list[i].guid
+            assert process.outputs[0].guid == ai_model.guid
 
 
 def test_ai_model_processes_creator(
@@ -217,10 +224,11 @@ def test_ai_model_processes_creator(
     )
     currnt_processes_sum = 0
     _assert_response_processes_creator(
-        mutation_response, list_training, AIDatasetType.TRAINING, 0, ai_model
+        client, mutation_response, list_training, AIDatasetType.TRAINING, 0, ai_model
     )
     currnt_processes_sum += len(list_training)
     _assert_response_processes_creator(
+        client,
         mutation_response,
         list_testing,
         AIDatasetType.TESTING,
@@ -229,6 +237,7 @@ def test_ai_model_processes_creator(
     )
     currnt_processes_sum += len(list_testing)
     _assert_response_processes_creator(
+        client,
         mutation_response,
         list_inference,
         AIDatasetType.INFERENCE,
@@ -237,6 +246,7 @@ def test_ai_model_processes_creator(
     )
     currnt_processes_sum += len(list_inference)
     _assert_response_processes_creator(
+        client,
         mutation_response,
         list_validation,
         AIDatasetType.VALIDATION,
@@ -245,6 +255,7 @@ def test_ai_model_processes_creator(
     )
     currnt_processes_sum += len(list_validation)
     _assert_response_processes_creator(
+        client,
         mutation_response,
         list_output,
         AIDatasetType.OUTPUT,

tests/integration/aio/test_admin.py

@@ -22,7 +22,7 @@
 
 EMAIL_DOMAIN = "@atlan.com"
 
-_default_group_count: int = 0
+_user_group_count_baseline: int = 0
 
 
 async def create_group_async(
@@ -44,9 +44,11 @@ async def test_retrieve_roles(client: AsyncAtlanClient):
 
 @pytest_asyncio.fixture(scope="module")
 async def group(client: AsyncAtlanClient) -> AsyncGenerator[CreateGroupResponse, None]:
+    global _user_group_count_baseline
     to_create = AtlanGroup.create(GROUP_NAME)
     fixed_user = await client.user.get_by_username(FIXED_USER)
     assert fixed_user
+    _user_group_count_baseline = fixed_user.group_count or 0
     g = await client.group.create(group=to_create, user_ids=[str(fixed_user.id)])
     yield g
     await delete_group_async(client, g.group)
@@ -89,13 +91,9 @@ async def test_create_group(client: AsyncAtlanClient, group: CreateGroupResponse
 async def test_retrieve_all_groups(
     client: AsyncAtlanClient, group: CreateGroupResponse
 ):
-    global _default_group_count
     groups = await client.group.get_all()
     assert groups.records
     assert len(groups.records) >= 1
-    for group1 in groups.records:
-        if group1.is_default():
-            _default_group_count += 1
 
 
 async def test_group_get_all_pagination(client: AsyncAtlanClient):
@@ -257,14 +255,13 @@ async def test_user_get_by_email_and_emails_pagination(client: AsyncAtlanClient)
 async def test_retrieve_existing_user(
     client: AsyncAtlanClient, group: CreateGroupResponse
 ):
-    global _default_group_count
     all_users = await client.user.get_all()
     assert all_users.records
     assert len(all_users.records) >= 1  # type: ignore
     user1 = await client.user.get_by_username(FIXED_USER)
     assert user1
     assert user1.id
-    assert user1.group_count == 1 + _default_group_count
+    assert user1.group_count == _user_group_count_baseline + 1
     response = await client.user.get_by_usernames(usernames=[FIXED_USER])
     assert response
     assert response.records is not None
@@ -351,16 +348,12 @@ async def test_final_user_state(
     client: AsyncAtlanClient,
     group: CreateGroupResponse,
 ):
-    global _default_group_count
     fixed_user = await client.user.get_by_username(FIXED_USER)
     assert fixed_user
     assert fixed_user.id
     response = await client.user.get_groups(fixed_user.id)
-    assert (
-        response.records is None
-        or len(response.records) == 0
-        or len(response.records) == _default_group_count
-    )
+    group_ids = {g.id for g in response.records or []}
+    assert group.group not in group_ids
 
 
 @pytest.mark.order(after="test_final_user_state")

tests/integration/aio/test_client.py

@@ -7,6 +7,13 @@
 import pytest_asyncio
 from httpx import Headers
 from pydantic.v1 import StrictStr
+from tenacity import (
+    RetryError,
+    retry,
+    retry_if_result,
+    stop_after_attempt,
+    wait_exponential,
+)
 
 from pyatlan import __version__ as VERSION
 from pyatlan.client.aio.client import AsyncAtlanClient
@@ -276,12 +283,22 @@ async def generate_audit_entries(
         time.sleep(1)
 
     request = AuditSearchRequest.by_guid(guid=audit_glossary.guid, size=log_count)
-    response = await client.audit.search(request)
-    assert response.total_count >= log_count, (
-        f"audit search failed, expected at least {log_count} log_count but got {response.total_count}."
+
+    @retry(
+        retry=retry_if_result(lambda count: count < log_count),
+        stop=stop_after_attempt(10),
+        wait=wait_exponential(multiplier=1, min=2, max=10),
+    )
+    async def _audit_entry_count() -> int:
+        return (await client.audit.search(request)).total_count
+
+    try:
+        total_count = await _audit_entry_count()
+    except RetryError as err:
+        total_count = err.last_attempt.result()
+    assert total_count >= log_count, (
+        f"audit search failed, expected at least {log_count} log_count but got {total_count}."
     )
-    # Force a wait to allow search entries to be indexed
-    time.sleep(10)
 
 
 async def _view_test_glossary_by_search(

tests/integration/test_client.py

@@ -6,6 +6,13 @@
 import pytest
 from httpx import Headers
 from pydantic.v1 import StrictStr
+from tenacity import (
+    RetryError,
+    retry,
+    retry_if_result,
+    stop_after_attempt,
+    wait_exponential,
+)
 
 from pyatlan import __version__ as VERSION
 from pyatlan.client.atlan import DEFAULT_RETRY, AtlanClient
@@ -1051,9 +1058,21 @@ def generate_audit_entries(client: AtlanClient, audit_glossary: AtlasGlossary):
         time.sleep(1)
 
     request = AuditSearchRequest.by_guid(guid=audit_glossary.guid, size=log_count)
-    response = client.audit.search(request)
-    assert response.total_count >= log_count, (
-        f"Expected at least {log_count} logs, but got {response.total_count}."
+
+    @retry(
+        retry=retry_if_result(lambda count: count < log_count),
+        stop=stop_after_attempt(10),
+        wait=wait_exponential(multiplier=1, min=2, max=10),
+    )
+    def _audit_entry_count() -> int:
+        return client.audit.search(request).total_count
+
+    try:
+        total_count = _audit_entry_count()
+    except RetryError as err:
+        total_count = err.last_attempt.result()
+    assert total_count >= log_count, (
+        f"Expected at least {log_count} logs, but got {total_count}."
     )