Skip to content

Commit 0a8cc3a

Browse files
committed
Added tests, fixed lint issue and introduced atlan error code
1 parent 857aa76 commit 0a8cc3a

7 files changed

Lines changed: 388 additions & 18 deletions

File tree

pyatlan/client/aio/oauth_client.py

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44
from __future__ import annotations
55

6-
from typing import TYPE_CHECKING, List, Optional
6+
from typing import List, Optional
77

88
from pydantic.v1 import validate_arguments
99

@@ -18,10 +18,11 @@
1818
RoleGet,
1919
)
2020
from pyatlan.errors import ErrorCode
21-
from pyatlan.model.oauth_clients import OAuthClient, OAuthClientCreateResponse, OAuthClientResponse
22-
23-
if TYPE_CHECKING:
24-
from pyatlan.client.aio import AsyncAtlanClient
21+
from pyatlan.model.oauth_clients import (
22+
OAuthClient,
23+
OAuthClientCreateResponse,
24+
OAuthClientResponse,
25+
)
2526

2627

2728
class AsyncOAuthClientClient:
@@ -147,7 +148,7 @@ async def create(
147148
:param persona_qns: qualified names of personas to associate with the OAuth client
148149
:returns: the created OAuthClientCreateResponse (includes client_id and client_secret)
149150
:raises AtlanError: on any API communication issue
150-
:raises ValueError: if the specified role description is not found
151+
:raises NotFoundError: if the specified role description is not found
151152
"""
152153
# Fetch available roles and resolve the user-provided role name
153154
available_roles = await self._fetch_available_roles()

pyatlan/client/common/oauth_client.py

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
GET_OAUTH_CLIENTS,
1313
UPDATE_OAUTH_CLIENT,
1414
)
15+
from pyatlan.errors import ErrorCode
1516
from pyatlan.model.oauth_clients import (
1617
OAuthClient,
1718
OAuthClientCreateResponse,
@@ -182,7 +183,7 @@ def resolve_role_name(role: str, available_roles: List[AtlanRole]) -> str:
182183
:param role: user-provided role description
183184
:param available_roles: list of available roles from the API
184185
:returns: the actual API role name (e.g., '$admin')
185-
:raises ValueError: if the role description is not found
186+
:raises NotFoundError: if the role description is not found
186187
"""
187188
role_lower = role.lower().strip()
188189

@@ -200,8 +201,8 @@ def resolve_role_name(role: str, available_roles: List[AtlanRole]) -> str:
200201
return desc_to_name[role_lower]
201202

202203
# No match found - raise error with available descriptions
203-
raise ValueError(
204-
f"Role '{role}' not found. Available roles: {', '.join(sorted(available_descriptions))}"
204+
raise ErrorCode.ROLE_NOT_FOUND_BY_DESCRIPTION.exception_with_parameters(
205+
role, ", ".join(sorted(available_descriptions))
205206
)
206207

207208
@staticmethod
@@ -213,9 +214,7 @@ def build_roles_filter() -> str:
213214
"""
214215
import json
215216

216-
return json.dumps(
217-
{"$or": [{"level": "workspace"}, {"level": "admin-subrole"}]}
218-
)
217+
return json.dumps({"$or": [{"level": "workspace"}, {"level": "admin-subrole"}]})
219218

220219
@staticmethod
221220
def prepare_request(

pyatlan/client/oauth_client.py

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
# Copyright 2025 Atlan Pte. Ltd.
33
from __future__ import annotations
44

5-
from typing import TYPE_CHECKING, List, Optional
5+
from typing import List, Optional
66

77
from pydantic.v1 import validate_arguments
88

@@ -17,10 +17,11 @@
1717
RoleGet,
1818
)
1919
from pyatlan.errors import ErrorCode
20-
from pyatlan.model.oauth_clients import OAuthClient, OAuthClientCreateResponse, OAuthClientResponse
21-
22-
if TYPE_CHECKING:
23-
from pyatlan.client.atlan import AtlanClient
20+
from pyatlan.model.oauth_clients import (
21+
OAuthClient,
22+
OAuthClientCreateResponse,
23+
OAuthClientResponse,
24+
)
2425

2526

2627
class OAuthClientClient:
@@ -146,7 +147,7 @@ def create(
146147
:param persona_qns: qualified names of personas to associate with the OAuth client
147148
:returns: the created OAuthClientCreateResponse (includes client_id and client_secret)
148149
:raises AtlanError: on any API communication issue
149-
:raises ValueError: if the specified role description is not found
150+
:raises NotFoundError: if the specified role description is not found
150151
"""
151152
# Fetch available roles and resolve the user-provided role name
152153
available_roles = self._fetch_available_roles()

pyatlan/errors.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -967,6 +967,13 @@ class ErrorCode(Enum):
967967
"Verify you have provided a valid DQ rule type.",
968968
NotFoundError,
969969
)
970+
ROLE_NOT_FOUND_BY_DESCRIPTION = (
971+
404,
972+
"ATLAN-PYTHON-404-030",
973+
"Role '{0}' does not exist. Available roles: {1}",
974+
"Verify the role description provided matches one of the available roles.",
975+
NotFoundError,
976+
)
970977
CONFLICT_PASSTHROUGH = (
971978
409,
972979
"ATLAN-PYTHON-409-000",

pyatlan/model/oauth_clients.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
from pyatlan.model.core import AtlanObject
88

9+
910
class OAuthClientRequest(AtlanObject):
1011
"""Request object for creating an OAuth client."""
1112

Lines changed: 181 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,181 @@
1+
# SPDX-License-Identifier: Apache-2.0
2+
# Copyright 2025 Atlan Pte. Ltd.
3+
"""Async integration tests for OAuth client CRUD operations."""
4+
5+
import time
6+
from typing import AsyncGenerator
7+
8+
import pytest
9+
import pytest_asyncio
10+
11+
from pyatlan.client.aio import AsyncAtlanClient
12+
from pyatlan.model.oauth_clients import OAuthClient, OAuthClientCreateResponse
13+
from tests.integration.client import TestId
14+
15+
MODULE_NAME = TestId.make_unique("AsyncOAuthClient")
16+
17+
# Test data
18+
OAUTH_CLIENT_NAME = f"{MODULE_NAME}_test_client"
19+
OAUTH_CLIENT_DESCRIPTION = "Async integration test OAuth client"
20+
OAUTH_CLIENT_DESCRIPTION_UPDATED = "Updated async integration test OAuth client"
21+
OAUTH_CLIENT_ROLE = "Admin" # Role description
22+
23+
24+
async def delete_oauth_client_async(client: AsyncAtlanClient, client_id: str) -> None:
25+
"""Helper to delete an OAuth client."""
26+
await client.oauth_client.purge(client_id)
27+
28+
29+
@pytest_asyncio.fixture(scope="module")
30+
async def oauth_client_response(
31+
client: AsyncAtlanClient,
32+
) -> AsyncGenerator[OAuthClientCreateResponse, None]:
33+
"""
34+
Fixture to create an OAuth client for testing.
35+
Yields the create response (which includes client_secret).
36+
Cleans up by deleting the OAuth client after tests.
37+
"""
38+
# Create OAuth client
39+
response = await client.oauth_client.create(
40+
name=OAUTH_CLIENT_NAME,
41+
role=OAUTH_CLIENT_ROLE,
42+
description=OAUTH_CLIENT_DESCRIPTION,
43+
)
44+
assert response is not None
45+
assert response.client_id is not None
46+
assert response.client_secret is not None
47+
48+
yield response
49+
50+
# Cleanup
51+
if response.client_id:
52+
await delete_oauth_client_async(client, response.client_id)
53+
54+
55+
def _assert_oauth_client_create_response(response: OAuthClientCreateResponse):
56+
"""Assert the OAuth client create response has expected values."""
57+
assert response is not None
58+
assert response.id is not None
59+
assert response.client_id is not None
60+
assert response.client_id.startswith("oauth-client-")
61+
assert response.client_secret is not None
62+
assert response.display_name == OAUTH_CLIENT_NAME
63+
assert response.description == OAUTH_CLIENT_DESCRIPTION
64+
assert response.created_by is not None
65+
assert response.created_at is not None
66+
67+
68+
def _assert_oauth_client(oauth_client: OAuthClient, is_updated: bool = False):
69+
"""Assert the OAuth client has expected values."""
70+
assert oauth_client is not None
71+
assert oauth_client.id is not None
72+
assert oauth_client.client_id is not None
73+
assert oauth_client.client_id.startswith("oauth-client-")
74+
assert oauth_client.display_name == OAUTH_CLIENT_NAME
75+
if is_updated:
76+
assert oauth_client.description == OAUTH_CLIENT_DESCRIPTION_UPDATED
77+
else:
78+
assert oauth_client.description == OAUTH_CLIENT_DESCRIPTION
79+
80+
81+
async def test_oauth_client_create(
82+
client: AsyncAtlanClient,
83+
oauth_client_response: OAuthClientCreateResponse,
84+
):
85+
"""Test creating an OAuth client."""
86+
_assert_oauth_client_create_response(oauth_client_response)
87+
88+
89+
@pytest.mark.order(after="test_oauth_client_create")
90+
async def test_oauth_client_get_by_id(
91+
client: AsyncAtlanClient,
92+
oauth_client_response: OAuthClientCreateResponse,
93+
):
94+
"""Test retrieving an OAuth client by ID."""
95+
assert oauth_client_response.client_id is not None
96+
time.sleep(2) # Allow time for eventual consistency
97+
98+
oauth_client = await client.oauth_client.get_by_id(oauth_client_response.client_id)
99+
_assert_oauth_client(oauth_client)
100+
101+
102+
@pytest.mark.order(after="test_oauth_client_get_by_id")
103+
async def test_oauth_client_get_all(
104+
client: AsyncAtlanClient,
105+
oauth_client_response: OAuthClientCreateResponse,
106+
):
107+
"""Test retrieving all OAuth clients."""
108+
assert oauth_client_response.client_id is not None
109+
time.sleep(2)
110+
111+
response = await client.oauth_client.get_all()
112+
assert response is not None
113+
assert response.records is not None
114+
assert len(response.records) >= 1
115+
116+
# Find our test client
117+
found = False
118+
for oauth_client in response.records:
119+
if oauth_client.client_id == oauth_client_response.client_id:
120+
found = True
121+
_assert_oauth_client(oauth_client)
122+
break
123+
assert found, f"OAuth client {oauth_client_response.client_id} not found"
124+
125+
126+
@pytest.mark.order(after="test_oauth_client_get_all")
127+
async def test_oauth_client_get_with_pagination(
128+
client: AsyncAtlanClient,
129+
oauth_client_response: OAuthClientCreateResponse,
130+
):
131+
"""Test retrieving OAuth clients with pagination."""
132+
assert oauth_client_response.client_id is not None
133+
134+
response = await client.oauth_client.get(limit=10, offset=0, sort="-createdAt")
135+
assert response is not None
136+
assert response.total_record is not None
137+
assert response.total_record >= 1
138+
139+
140+
@pytest.mark.order(after="test_oauth_client_get_with_pagination")
141+
async def test_oauth_client_update_description(
142+
client: AsyncAtlanClient,
143+
oauth_client_response: OAuthClientCreateResponse,
144+
):
145+
"""Test updating an OAuth client's description."""
146+
assert oauth_client_response.client_id is not None
147+
time.sleep(2)
148+
149+
updated = await client.oauth_client.update(
150+
client_id=oauth_client_response.client_id,
151+
description=OAUTH_CLIENT_DESCRIPTION_UPDATED,
152+
)
153+
_assert_oauth_client(updated, is_updated=True)
154+
155+
156+
@pytest.mark.order(after="test_oauth_client_update_description")
157+
async def test_oauth_client_verify_update_persisted(
158+
client: AsyncAtlanClient,
159+
oauth_client_response: OAuthClientCreateResponse,
160+
):
161+
"""Verify that the update was persisted."""
162+
assert oauth_client_response.client_id is not None
163+
time.sleep(2)
164+
165+
oauth_client = await client.oauth_client.get_by_id(oauth_client_response.client_id)
166+
_assert_oauth_client(oauth_client, is_updated=True)
167+
168+
169+
async def test_oauth_client_create_with_invalid_role_raises_error(
170+
client: AsyncAtlanClient,
171+
):
172+
"""Test that creating an OAuth client with an invalid role raises an error."""
173+
from pyatlan.errors import NotFoundError
174+
175+
with pytest.raises(NotFoundError) as exc_info:
176+
await client.oauth_client.create(
177+
name="test-invalid-role",
178+
role="InvalidRole",
179+
)
180+
assert "does not exist" in str(exc_info.value)
181+
assert "Available roles:" in str(exc_info.value)

0 commit comments

Comments
 (0)