diff --git a/tests/integration/admin_test.py b/tests/integration/admin_test.py index 397574399..23e0ff68e 100644 --- a/tests/integration/admin_test.py +++ b/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") diff --git a/tests/integration/ai_asset_test.py b/tests/integration/ai_asset_test.py index bf8d2a8e5..a5d469f96 100644 --- a/tests/integration/ai_asset_test.py +++ b/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, diff --git a/tests/integration/aio/test_admin.py b/tests/integration/aio/test_admin.py index cd28f79dd..2324c476e 100644 --- a/tests/integration/aio/test_admin.py +++ b/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") diff --git a/tests/integration/aio/test_client.py b/tests/integration/aio/test_client.py index 705ed3188..5e2a5b8e0 100644 --- a/tests/integration/aio/test_client.py +++ b/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( diff --git a/tests/integration/test_client.py b/tests/integration/test_client.py index a27fc74f2..3e5eca5c3 100644 --- a/tests/integration/test_client.py +++ b/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}." )