[fix] Fixed open_lineage.create_connection() return

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(

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)