Skip to content

Commit 442b375

Browse files
Aryamanz29Aryamanz29
committed
[fix] Fixed connection asset -> caches imports
1 parent db9186a commit 442b375

File tree

5 files changed

+128
-22
lines changed

5 files changed

+128
-22
lines changed

‎pyatlan/generator/templates/methods/asset/purpose.jinja2

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -51,10 +51,12 @@
5151
policy.policy_groups = {"public"}
5252
else:
5353
if policy_groups:
54-
from pyatlan.cache.group_cache import GroupCache
54+
from pyatlan.client.atlan import AtlanClient
5555

5656
for group_name in policy_groups:
57-
if not GroupCache.get_id_for_name(group_name):
57+
if not AtlanClient.get_current_client().group_cache.get_id_for_name(
58+
group_name
59+
):
5860
raise ValueError(
5961
f"Provided group name {group_name} was not found in Atlan."
6062
)
@@ -63,10 +65,12 @@
6365
else:
6466
policy.policy_groups = None
6567
if policy_users:
66-
from pyatlan.cache.user_cache import UserCache
68+
from pyatlan.client.atlan import AtlanClient
6769

6870
for username in policy_users:
69-
if not UserCache.get_id_for_name(username):
71+
if not AtlanClient.get_current_client().user_cache.get_id_for_name(
72+
username
73+
):
7074
raise ValueError(
7175
f"Provided username {username} was not found in Atlan."
7276
)
@@ -108,10 +112,12 @@
108112
policy.policy_groups = {"public"}
109113
else:
110114
if policy_groups:
111-
from pyatlan.cache.group_cache import GroupCache
115+
from pyatlan.client.atlan import AtlanClient
112116

113117
for group_name in policy_groups:
114-
if not GroupCache.get_id_for_name(group_name):
118+
if not AtlanClient.get_current_client().group_cache.get_id_for_name(
119+
group_name
120+
):
115121
raise ValueError(
116122
f"Provided group name {group_name} was not found in Atlan."
117123
)
@@ -120,10 +126,12 @@
120126
else:
121127
policy.policy_groups = None
122128
if policy_users:
123-
from pyatlan.cache.user_cache import UserCache
129+
from pyatlan.client.atlan import AtlanClient
124130

125131
for username in policy_users:
126-
if not UserCache.get_id_for_name(username):
132+
if not AtlanClient.get_current_client().user_cache.get_id_for_name(
133+
username
134+
):
127135
raise ValueError(
128136
f"Provided username {username} was not found in Atlan."
129137
)

‎pyatlan/generator/templates/methods/attribute/connection.jinja2

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

44
@validator("admin_users")
55
def admin_users_valid(cls, admin_users, values):
6-
from pyatlan.cache.user_cache import UserCache
6+
from pyatlan.client.atlan import AtlanClient
7+
78
if values.get("is_loaded", False):
8-
UserCache.validate_names(names=admin_users)
9+
AtlanClient.get_current_client().user_cache.validate_names(
10+
names=admin_users
11+
)
912
return admin_users
1013

1114
@validator("admin_roles")
1215
def admin_roles_valid(cls, admin_roles, values):
13-
from pyatlan.cache.role_cache import RoleCache
16+
from pyatlan.client.atlan import AtlanClient
17+
1418
if values.get("is_loaded", False):
15-
RoleCache.validate_idstrs(idstrs=admin_roles)
19+
AtlanClient.get_current_client().role_cache.validate_idstrs(
20+
idstrs=admin_roles
21+
)
1622
return admin_roles
1723

1824
@validator("admin_groups")
1925
def admin_groups_valid(cls, admin_groups, values):
20-
from pyatlan.cache.group_cache import GroupCache
26+
from pyatlan.client.atlan import AtlanClient
27+
2128
if values.get("is_loaded", False):
22-
GroupCache.validate_aliases(aliases=admin_groups)
29+
AtlanClient.get_current_client().group_cache.validate_aliases(
30+
aliases=admin_groups
31+
)
2332
return admin_groups

‎pyatlan/model/assets/connection.py

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -670,26 +670,32 @@ class Attributes(Asset.Attributes):
670670

671671
@validator("admin_users")
672672
def admin_users_valid(cls, admin_users, values):
673-
from pyatlan.cache.user_cache import UserCache
673+
from pyatlan.client.atlan import AtlanClient
674674

675675
if values.get("is_loaded", False):
676-
UserCache.validate_names(names=admin_users)
676+
AtlanClient.get_current_client().user_cache.validate_names(
677+
names=admin_users
678+
)
677679
return admin_users
678680

679681
@validator("admin_roles")
680682
def admin_roles_valid(cls, admin_roles, values):
681-
from pyatlan.cache.role_cache import RoleCache
683+
from pyatlan.client.atlan import AtlanClient
682684

683685
if values.get("is_loaded", False):
684-
RoleCache.validate_idstrs(idstrs=admin_roles)
686+
AtlanClient.get_current_client().role_cache.validate_idstrs(
687+
idstrs=admin_roles
688+
)
685689
return admin_roles
686690

687691
@validator("admin_groups")
688692
def admin_groups_valid(cls, admin_groups, values):
689-
from pyatlan.cache.group_cache import GroupCache
693+
from pyatlan.client.atlan import AtlanClient
690694

691695
if values.get("is_loaded", False):
692-
GroupCache.validate_aliases(aliases=admin_groups)
696+
AtlanClient.get_current_client().group_cache.validate_aliases(
697+
aliases=admin_groups
698+
)
693699
return admin_groups
694700

695701
attributes: Connection.Attributes = Field(

‎tests/unit/model/connection_test.py

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
from typing import List, Optional
2-
from unittest.mock import patch
2+
from unittest.mock import MagicMock, patch
33

44
import pytest
55

@@ -10,6 +10,47 @@
1010
from tests.unit.model.constants import CONNECTION_NAME, CONNECTION_QUALIFIED_NAME
1111

1212

13+
@pytest.fixture(autouse=True)
14+
def set_env(monkeypatch):
15+
monkeypatch.setenv("ATLAN_BASE_URL", "https://test.atlan.com")
16+
monkeypatch.setenv("ATLAN_API_KEY", "test-api-key")
17+
18+
19+
@pytest.fixture()
20+
def client():
21+
return AtlanClient()
22+
23+
24+
@pytest.fixture()
25+
def current_client(client, monkeypatch):
26+
monkeypatch.setattr(
27+
AtlanClient,
28+
"get_current_client",
29+
lambda: client,
30+
)
31+
32+
33+
@pytest.fixture()
34+
def mock_group_cache(current_client, monkeypatch):
35+
mock_cache = MagicMock()
36+
monkeypatch.setattr(AtlanClient, "group_cache", mock_cache)
37+
return mock_cache
38+
39+
40+
@pytest.fixture()
41+
def mock_user_cache(current_client, monkeypatch):
42+
mock_cache = MagicMock()
43+
monkeypatch.setattr(AtlanClient, "user_cache", mock_cache)
44+
return mock_cache
45+
46+
47+
@pytest.fixture()
48+
def mock_role_cache(current_client, monkeypatch):
49+
mock_cache = MagicMock()
50+
monkeypatch.setattr(AtlanClient, "role_cache", mock_cache)
51+
return mock_cache
52+
53+
1354
@pytest.mark.parametrize(
1455
"name, connector_type, admin_users, admin_groups, admin_roles, message",
1556
[

‎tests/unit/test_packages.py

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
from json import load, loads
22
from pathlib import Path
3-
from unittest.mock import patch
3+
from unittest.mock import MagicMock, patch
44

55
import pytest
66

7+
from pyatlan.client.atlan import AtlanClient
78
from pyatlan.errors import InvalidRequestError
89
from pyatlan.model.assets.core import Asset
910
from pyatlan.model.enums import AssetDeltaHandling, AssetInputHandling, AssetRemovalType
@@ -130,6 +131,47 @@ def mock_connection_guid():
130131
yield mock_random
131132

132133

134+
@pytest.fixture(autouse=True)
135+
def set_env(monkeypatch):
136+
monkeypatch.setenv("ATLAN_BASE_URL", "https://test.atlan.com")
137+
monkeypatch.setenv("ATLAN_API_KEY", "test-api-key")
138+
139+
140+
@pytest.fixture()
141+
def client():
142+
return AtlanClient()
143+
144+
145+
@pytest.fixture()
146+
def current_client(client, monkeypatch):
147+
monkeypatch.setattr(
148+
AtlanClient,
149+
"get_current_client",
150+
lambda: client,
151+
)
152+
153+
154+
@pytest.fixture()
155+
def mock_group_cache(current_client, monkeypatch):
156+
mock_cache = MagicMock()
157+
monkeypatch.setattr(AtlanClient, "group_cache", mock_cache)
158+
return mock_cache
159+
160+
161+
@pytest.fixture()
162+
def mock_user_cache(current_client, monkeypatch):
163+
mock_cache = MagicMock()
164+
monkeypatch.setattr(AtlanClient, "user_cache", mock_cache)
165+
return mock_cache
166+
167+
168+
@pytest.fixture()
169+
def mock_role_cache(current_client, monkeypatch):
170+
mock_cache = MagicMock()
171+
monkeypatch.setattr(AtlanClient, "role_cache", mock_cache)
172+
return mock_cache
173+
174+
133175
@pytest.fixture()
134176
def mock_package_env(
135177
mock_role_cache,

0 commit comments

Comments
 (0)