From bda3bcf33c9088739d28db5d8f2f5aaaea4a7b00 Mon Sep 17 00:00:00 2001 From: Aryamanz29 Date: Fri, 17 Jan 2025 16:49:24 +0530 Subject: [PATCH] [fix] Fixed `open_lineage.create_connection()` return --- pyatlan/client/open_lineage.py | 6 +++--- tests/integration/test_open_lineage.py | 15 ++++++++------- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/pyatlan/client/open_lineage.py b/pyatlan/client/open_lineage.py index db32c59a..48d22621 100644 --- a/pyatlan/client/open_lineage.py +++ b/pyatlan/client/open_lineage.py @@ -11,6 +11,7 @@ from pyatlan.model.credential import CredentialResponse from pyatlan.model.enums import AtlanConnectorType from pyatlan.model.open_lineage.event import OpenLineageEvent +from pyatlan.model.response import AssetMutationResponse class OpenLineageClient: @@ -60,7 +61,7 @@ def create_connection( admin_users: Optional[List[str]] = None, admin_roles: Optional[List[str]] = None, admin_groups: Optional[List[str]] = None, - ) -> Connection: + ) -> AssetMutationResponse: """ Creates a connection for OpenLineage. @@ -85,8 +86,7 @@ def create_connection( ) connection.default_credential_guid = response.id - response = client.asset.save(connection) # type: ignore - return connection + return client.asset.save(connection) @validate_arguments def send( diff --git a/tests/integration/test_open_lineage.py b/tests/integration/test_open_lineage.py index b64dd0b9..f4ff78fa 100644 --- a/tests/integration/test_open_lineage.py +++ b/tests/integration/test_open_lineage.py @@ -17,14 +17,16 @@ @pytest.fixture(scope="module") def connection(client: AtlanClient): - admin_role_guid = str(RoleCache.get_id_for_name("$admin")) - - c = client.open_lineage.create_connection( + admin_role_guid = RoleCache.get_id_for_name("$admin") + assert admin_role_guid + response = client.open_lineage.create_connection( name=MODULE_NAME, admin_roles=[admin_role_guid] ) - guid = c.guid - yield c - delete_asset(client, asset_type=Connection, guid=guid) + result = response.assets_created(asset_type=Connection)[0] + yield client.asset.get_by_guid( + result.guid, asset_type=Connection, ignore_relationships=False + ) + delete_asset(client, asset_type=Connection, guid=result.guid) def test_open_lineage_integration(connection: Connection, client: AtlanClient): @@ -123,6 +125,5 @@ def test_open_lineage_integration(connection: Connection, client: AtlanClient): process.get("uniqueAttributes", {}).get("qualifiedName") == f"{connection.qualified_name}/dag_123/process" ) - delete_asset(client, asset_type=Process, guid=process.get("guid")) delete_asset(client, asset_type=SparkJob, guid=job_asset.detail.guid)