Skip to content

Commit 612049f

Browse files
ErnestoLomaErnestoLoma
authoredMay 4, 2023
Merge pull request #43 from atlanhq/link-terms
Update code to link terms to assets
·
6.0.00.0.26
2 parents e94ed3f + fff3d1d commit 612049f

File tree

8 files changed

+3080
-47
lines changed

8 files changed

+3080
-47
lines changed
 

‎HISTORY.md

+6
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,9 @@
1+
## 0.0.26 (May 4, 2023)
2+
3+
* Add remove_terms method to AtlanClient
4+
* Add append_terms method to AtlanClient
5+
* Add replace_terms method to AtlanClient
6+
17
## 0.0.25 (May 2, 2023)
28

39
* Update create method for Readme asset

‎pyatlan/client/atlan.py

+112-12
Original file line numberDiff line numberDiff line change
@@ -313,8 +313,7 @@ def get_asset_by_qualified_name(
313313
raw_json["entity"]["attributes"].update(
314314
raw_json["entity"]["relationshipAttributes"]
315315
)
316-
asset = AssetResponse[A](**raw_json).entity
317-
asset.is_incomplete = False
316+
asset = self.handle_relationships(raw_json)
318317
if not isinstance(asset, asset_type):
319318
raise NotFoundError(
320319
message=f"Asset with qualifiedName {qualified_name} "
@@ -345,16 +344,7 @@ def get_asset_by_guid(
345344
GET_ENTITY_BY_GUID.format_path_with_params(guid),
346345
query_params,
347346
)
348-
if (
349-
"relationshipAttributes" in raw_json["entity"]
350-
and raw_json["entity"]["relationshipAttributes"]
351-
):
352-
raw_json["entity"]["attributes"].update(
353-
raw_json["entity"]["relationshipAttributes"]
354-
)
355-
raw_json["entity"]["relationshipAttributes"] = {}
356-
asset = AssetResponse[A](**raw_json).entity
357-
asset.is_incomplete = False
347+
asset = self.handle_relationships(raw_json)
358348
if not isinstance(asset, asset_type):
359349
raise NotFoundError(
360350
message=f"Asset with GUID {guid} is not of the type requested: {asset_type.__name__}.",
@@ -366,6 +356,19 @@ def get_asset_by_guid(
366356
raise NotFoundError(message=ae.user_message, code=ae.code) from ae
367357
raise ae
368358

359+
def handle_relationships(self, raw_json):
360+
if (
361+
"relationshipAttributes" in raw_json["entity"]
362+
and raw_json["entity"]["relationshipAttributes"]
363+
):
364+
raw_json["entity"]["attributes"].update(
365+
raw_json["entity"]["relationshipAttributes"]
366+
)
367+
raw_json["entity"]["relationshipAttributes"] = {}
368+
asset = AssetResponse[A](**raw_json).entity
369+
asset.is_incomplete = False
370+
return asset
371+
369372
@validate_arguments()
370373
def retrieve_minimal(self, guid: str, asset_type: Type[A]) -> A:
371374
return self.get_asset_by_guid(
@@ -602,3 +605,100 @@ def replace_custom_metadata(self, guid: str, custom_metadata: CustomMetadata):
602605
None,
603606
custom_metadata_request,
604607
)
608+
609+
@validate_arguments()
610+
def append_terms(
611+
self,
612+
asset_type: Type[A],
613+
terms: list[AtlasGlossaryTerm],
614+
guid: Optional[str] = None,
615+
qualified_name: Optional[str] = None,
616+
) -> A:
617+
if guid:
618+
if qualified_name:
619+
raise ValueError(
620+
"Either guid or qualified_name can be be specified not both"
621+
)
622+
asset = self.get_asset_by_guid(guid=guid, asset_type=asset_type)
623+
elif qualified_name:
624+
asset = self.get_asset_by_qualified_name(
625+
qualified_name=qualified_name, asset_type=asset_type
626+
)
627+
else:
628+
raise ValueError("Either guid or qualified name must be specified")
629+
if not terms:
630+
return asset
631+
replacement_terms: list[AtlasGlossaryTerm] = []
632+
if existing_terms := asset.terms:
633+
replacement_terms.extend(
634+
term for term in existing_terms if term.relationship_status != "DELETED"
635+
)
636+
replacement_terms.extend(terms)
637+
asset.terms = replacement_terms
638+
response = self.upsert(entity=asset)
639+
if assets := response.assets_updated(asset_type=asset_type):
640+
return assets[0]
641+
return asset
642+
643+
@validate_arguments()
644+
def replace_terms(
645+
self,
646+
asset_type: Type[A],
647+
terms: list[AtlasGlossaryTerm],
648+
guid: Optional[str] = None,
649+
qualified_name: Optional[str] = None,
650+
) -> A:
651+
if guid:
652+
if qualified_name:
653+
raise ValueError(
654+
"Either guid or qualified_name can be be specified not both"
655+
)
656+
asset = self.get_asset_by_guid(guid=guid, asset_type=asset_type)
657+
elif qualified_name:
658+
asset = self.get_asset_by_qualified_name(
659+
qualified_name=qualified_name, asset_type=asset_type
660+
)
661+
else:
662+
raise ValueError("Either guid or qualified name must be specified")
663+
asset.terms = terms
664+
response = self.upsert(entity=asset)
665+
if assets := response.assets_updated(asset_type=asset_type):
666+
return assets[0]
667+
return asset
668+
669+
@validate_arguments()
670+
def remove_terms(
671+
self,
672+
asset_type: Type[A],
673+
terms: list[AtlasGlossaryTerm],
674+
guid: Optional[str] = None,
675+
qualified_name: Optional[str] = None,
676+
) -> A:
677+
if not terms:
678+
raise ValueError("A list of terms to remove must be specified")
679+
if guid:
680+
if qualified_name:
681+
raise ValueError(
682+
"Either guid or qualified_name can be be specified not both"
683+
)
684+
asset = self.get_asset_by_guid(guid=guid, asset_type=asset_type)
685+
elif qualified_name:
686+
asset = self.get_asset_by_qualified_name(
687+
qualified_name=qualified_name, asset_type=asset_type
688+
)
689+
else:
690+
raise ValueError("Either guid or qualified name must be specified")
691+
replacement_terms: list[AtlasGlossaryTerm] = []
692+
guids_to_be_removed = {t.guid for t in terms}
693+
if existing_terms := asset.terms:
694+
replacement_terms.extend(
695+
term
696+
for term in existing_terms
697+
if term.relationship_status != "DELETED"
698+
and term.guid not in guids_to_be_removed
699+
)
700+
asset.terms = replacement_terms
701+
response = self.upsert(entity=asset)
702+
if assets := response.assets_updated(asset_type=asset_type):
703+
return assets[0]
704+
return asset

‎pyatlan/generator/templates/entity.jinja2

+19-3
Original file line numberDiff line numberDiff line change
@@ -71,11 +71,15 @@ def validate_required_fields(field_names:list[str], values:list[Any]):
7171
raise ValueError(f"{field_name} is required")
7272
if isinstance(value, str) and not value.strip():
7373
raise ValueError(f"{field_name} cannot be blank")
74-
{%- macro gen_properties(attribute_defs) %}
74+
{%- macro gen_properties(attribute_defs, additional_names=[]) %}
7575
_convience_properties: ClassVar[list[str]] = [
7676
{%- for attribute_def in attribute_defs %}
7777
"{{ attribute_def.name | to_snake_case }}",
78-
{%- endfor %}]
78+
{%- endfor %}
79+
{%- for name in additional_names %}
80+
"{{ name }}",
81+
{%- endfor %}]
82+
7983
{%- for attribute_def in attribute_defs %}
8084
{%- set type = attribute_def.typeName | get_type %}
8185
{%- set property_type %}{% if attribute_def.isOptional %}Optional[{% endif %}{{type}}{% if attribute_def.isOptional %}]{% endif %}{% endset %}
@@ -107,7 +111,19 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
107111
return object.__setattr__(self, name, value)
108112
super().__setattr__( name, value)
109113

110-
{{ gen_properties(entity_def.attribute_defs) }}
114+
{{ gen_properties(entity_def.attribute_defs, ["terms"]) }}
115+
116+
@property
117+
def terms(self) -> list[AtlasGlossaryTerm]:
118+
if self.attributes is None:
119+
self.attributes = self.Attributes()
120+
return [] if self.attributes.meanings is None else self.attributes.meanings
121+
122+
@terms.setter
123+
def terms(self, terms: list[AtlasGlossaryTerm]):
124+
if self.attributes is None:
125+
self.attributes = self.Attributes()
126+
self.attributes.meanings = terms
111127

112128
{%- if entity_def.name == "Referenceable" %}
113129

‎pyatlan/model/assets.py

+2,419-30
Large diffs are not rendered by default.

‎pyatlan/version.txt

+1-1
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
0.0.25
1+
0.0.26

‎tests/integration/test_client.py

+206
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,66 @@
1+
from itertools import count
2+
from typing import Callable, Generator
3+
14
import pytest
25

36
from pyatlan.client.atlan import AtlanClient
7+
from pyatlan.model.assets import AtlasGlossary, AtlasGlossaryTerm, Connection, Database
8+
9+
iter_count = count(1)
10+
11+
12+
@pytest.fixture(scope="module")
13+
def client() -> AtlanClient:
14+
return AtlanClient()
15+
16+
17+
@pytest.fixture(scope="module")
18+
def connection(client: AtlanClient) -> Connection:
19+
return client.get_asset_by_guid("b3a5c49a-0c7c-4e66-8453-f4da8d9ce222", Connection)
20+
21+
22+
@pytest.fixture(scope="module")
23+
def glossary(client: AtlanClient) -> Generator[AtlasGlossary, None, None]:
24+
glossary = AtlasGlossary.create(name="Integration Test Glossary")
25+
glossary = client.upsert(glossary).assets_created(asset_type=AtlasGlossary)[0]
26+
yield glossary
27+
client.purge_entity_by_guid(guid=glossary.guid)
28+
29+
30+
@pytest.fixture()
31+
def database(
32+
client: AtlanClient, connection: Connection
33+
) -> Generator[Database, None, None]:
34+
35+
database = Database.create(
36+
name=f"Integration_Test_Entity_DB{next(iter_count)}",
37+
connection_qualified_name=connection.attributes.qualified_name,
38+
)
39+
database = client.upsert(database).assets_created(Database)[0]
40+
41+
yield database
42+
43+
client.purge_entity_by_guid(guid=database.guid)
44+
45+
46+
@pytest.fixture()
47+
def make_term(
48+
client: AtlanClient, glossary
49+
) -> Generator[Callable[[str], AtlasGlossaryTerm], None, None]:
50+
created_term_guids = []
51+
52+
def _make_term(name: str) -> AtlasGlossaryTerm:
53+
term = AtlasGlossaryTerm.create(
54+
name=f"Integration Test Glossary Term {name}", anchor=glossary
55+
)
56+
term = client.upsert(term).assets_created(AtlasGlossaryTerm)[0]
57+
created_term_guids.append(term.guid)
58+
return term
59+
60+
yield _make_term
61+
62+
for guid in created_term_guids:
63+
client.purge_entity_by_guid(guid=guid)
464

565

666
def test_register_client_with_bad_parameter_raises_valueerror():
@@ -13,3 +73,149 @@ def test_register_client():
1373
client = AtlanClient(base_url="http://mark.atlan.com", api_key="123")
1474
AtlanClient.register_client(client)
1575
assert AtlanClient.get_default_client() == client
76+
77+
78+
def test_append_terms_with_guid(
79+
client: AtlanClient,
80+
make_term: Callable[[str], AtlasGlossaryTerm],
81+
database: Database,
82+
):
83+
term = make_term("Term1")
84+
85+
assert (
86+
database := client.append_terms(
87+
guid=database.guid, asset_type=Database, terms=[term]
88+
)
89+
)
90+
database = client.get_asset_by_guid(guid=database.guid, asset_type=Database)
91+
assert len(database.terms) == 1
92+
assert database.terms[0].guid == term.guid
93+
94+
95+
def test_append_terms_with_qualified_name(
96+
client: AtlanClient,
97+
make_term: Callable[[str], AtlasGlossaryTerm],
98+
database: Database,
99+
):
100+
term = make_term("Term1")
101+
102+
assert (
103+
database := client.append_terms(
104+
qualified_name=database.qualified_name, asset_type=Database, terms=[term]
105+
)
106+
)
107+
database = client.get_asset_by_guid(guid=database.guid, asset_type=Database)
108+
assert len(database.terms) == 1
109+
assert database.terms[0].guid == term.guid
110+
111+
112+
def test_append_terms_using_ref_by_guid_for_term(
113+
client: AtlanClient,
114+
make_term: Callable[[str], AtlasGlossaryTerm],
115+
database: Database,
116+
):
117+
term = make_term("Term1")
118+
119+
assert (
120+
database := client.append_terms(
121+
qualified_name=database.qualified_name,
122+
asset_type=Database,
123+
terms=[AtlasGlossaryTerm.ref_by_guid(guid=term.guid)],
124+
)
125+
)
126+
database = client.get_asset_by_guid(guid=database.guid, asset_type=Database)
127+
assert len(database.terms) == 1
128+
assert database.terms[0].guid == term.guid
129+
130+
131+
def test_replace_a_term(
132+
client: AtlanClient,
133+
make_term: Callable[[str], AtlasGlossaryTerm],
134+
database: Database,
135+
):
136+
original_term = make_term("Term1")
137+
assert (
138+
database := client.append_terms(
139+
qualified_name=database.qualified_name,
140+
asset_type=Database,
141+
terms=[AtlasGlossaryTerm.ref_by_guid(guid=original_term.guid)],
142+
)
143+
)
144+
145+
replacemant_term = make_term("Term2")
146+
assert (
147+
database := client.replace_terms(
148+
guid=database.guid, asset_type=Database, terms=[replacemant_term]
149+
)
150+
)
151+
152+
database = client.get_asset_by_guid(guid=database.guid, asset_type=Database)
153+
assert len(database.terms) == 2
154+
deleted_terms = [t for t in database.terms if t.relationship_status == "DELETED"]
155+
assert len(deleted_terms) == 1
156+
assert deleted_terms[0].guid == original_term.guid
157+
active_terms = [t for t in database.terms if t.relationship_status != "DELETED"]
158+
assert len(active_terms) == 1
159+
assert active_terms[0].guid == replacemant_term.guid
160+
161+
162+
def test_replace_all_term(
163+
client: AtlanClient,
164+
make_term: Callable[[str], AtlasGlossaryTerm],
165+
database: Database,
166+
):
167+
original_term = make_term("Term1")
168+
assert (
169+
database := client.append_terms(
170+
qualified_name=database.qualified_name,
171+
asset_type=Database,
172+
terms=[AtlasGlossaryTerm.ref_by_guid(guid=original_term.guid)],
173+
)
174+
)
175+
176+
assert (
177+
database := client.replace_terms(
178+
guid=database.guid, asset_type=Database, terms=[]
179+
)
180+
)
181+
182+
database = client.get_asset_by_guid(guid=database.guid, asset_type=Database)
183+
assert len(database.terms) == 1
184+
deleted_terms = [t for t in database.terms if t.relationship_status == "DELETED"]
185+
assert len(deleted_terms) == 1
186+
assert deleted_terms[0].guid == original_term.guid
187+
188+
189+
def test_remove_term(
190+
client: AtlanClient,
191+
make_term: Callable[[str], AtlasGlossaryTerm],
192+
database: Database,
193+
):
194+
original_term = make_term("Term1")
195+
another_term = make_term("Term2")
196+
assert (
197+
database := client.append_terms(
198+
qualified_name=database.qualified_name,
199+
asset_type=Database,
200+
terms=[
201+
AtlasGlossaryTerm.ref_by_guid(guid=original_term.guid),
202+
AtlasGlossaryTerm.ref_by_guid(guid=another_term.guid),
203+
],
204+
)
205+
)
206+
207+
assert (
208+
database := client.remove_terms(
209+
guid=database.guid,
210+
asset_type=Database,
211+
terms=[AtlasGlossaryTerm.ref_by_guid(original_term.guid)],
212+
)
213+
)
214+
215+
database = client.get_asset_by_guid(guid=database.guid, asset_type=Database)
216+
assert len(database.terms) == 2
217+
deleted_terms = [t for t in database.terms if t.relationship_status == "DELETED"]
218+
assert len(deleted_terms) == 1
219+
assert deleted_terms[0].guid == original_term.guid
220+
active_terms = [t for t in database.terms if t.relationship_status != "DELETED"]
221+
assert active_terms[0].guid == another_term.guid

‎tests/unit/test_client.py

+315
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,315 @@
1+
# SPDX-License-Identifier: Apache-2.0
2+
# Copyright 2022 Atlan Pte. Ltd.
3+
from unittest.mock import DEFAULT, patch
4+
5+
import pytest
6+
7+
from pyatlan.client.atlan import AtlanClient
8+
from pyatlan.model.assets import AtlasGlossaryTerm, Table
9+
10+
11+
@pytest.mark.parametrize(
12+
"guid, qualified_name, asset_type, terms, message",
13+
[
14+
(
15+
"123",
16+
None,
17+
Table,
18+
None,
19+
"1 validation error for AppendTerms\\nterms\\n none is not an allowed value ",
20+
),
21+
(
22+
None,
23+
None,
24+
Table,
25+
[AtlasGlossaryTerm()],
26+
"Either guid or qualified name must be specified",
27+
),
28+
(
29+
"123",
30+
None,
31+
None,
32+
[AtlasGlossaryTerm()],
33+
"1 validation error for AppendTerms\\nasset_type\\n none is not an allowed value ",
34+
),
35+
(
36+
"123",
37+
"default/abc",
38+
Table,
39+
[AtlasGlossaryTerm()],
40+
"Either guid or qualified_name can be be specified not both",
41+
),
42+
],
43+
)
44+
def test_append_terms_with_invalid_parameter_raises_valueerror(
45+
guid,
46+
qualified_name,
47+
asset_type,
48+
terms,
49+
message,
50+
monkeypatch,
51+
):
52+
monkeypatch.setenv("ATLAN_BASE_URL", "https://name.atlan.com")
53+
monkeypatch.setenv("ATLAN_API_KEY", "abkj")
54+
client = AtlanClient()
55+
56+
with pytest.raises(ValueError, match=message):
57+
client.append_terms(
58+
guid=guid, qualified_name=qualified_name, asset_type=asset_type, terms=terms
59+
)
60+
61+
62+
def test_append_with_valid_guid_and_no_terms_returns_asset(monkeypatch):
63+
monkeypatch.setenv("ATLAN_BASE_URL", "https://name.atlan.com")
64+
monkeypatch.setenv("ATLAN_API_KEY", "abkj")
65+
asset_type = Table
66+
table = asset_type()
67+
68+
with patch.object(
69+
AtlanClient, "get_asset_by_guid", return_value=table
70+
) as mock_method:
71+
client = AtlanClient()
72+
guid = "123"
73+
terms = []
74+
75+
assert (
76+
client.append_terms(guid=guid, asset_type=asset_type, terms=terms) == table
77+
)
78+
mock_method.assert_called_once_with(guid=guid, asset_type=asset_type)
79+
80+
81+
def test_append_with_valid_guid_when_no_terms_present_returns_asset_with_given_terms(
82+
monkeypatch,
83+
):
84+
monkeypatch.setenv("ATLAN_BASE_URL", "https://name.atlan.com")
85+
monkeypatch.setenv("ATLAN_API_KEY", "abkj")
86+
asset_type = Table
87+
with patch.multiple(
88+
AtlanClient, get_asset_by_guid=DEFAULT, upsert=DEFAULT
89+
) as mock_methods:
90+
table = Table()
91+
mock_methods["get_asset_by_guid"].return_value = table
92+
mock_methods["upsert"].return_value.assets_updated.return_value = [table]
93+
client = AtlanClient()
94+
guid = "123"
95+
terms = [AtlasGlossaryTerm()]
96+
97+
assert (
98+
asset := client.append_terms(guid=guid, asset_type=asset_type, terms=terms)
99+
)
100+
assert asset.terms == terms
101+
102+
103+
def test_append_with_valid_guid_when_deleted_terms_present_returns_asset_with_given_terms(
104+
monkeypatch,
105+
):
106+
monkeypatch.setenv("ATLAN_BASE_URL", "https://name.atlan.com")
107+
monkeypatch.setenv("ATLAN_API_KEY", "abkj")
108+
asset_type = Table
109+
with patch.multiple(
110+
AtlanClient, get_asset_by_guid=DEFAULT, upsert=DEFAULT
111+
) as mock_methods:
112+
table = Table(attributes=Table.Attributes())
113+
term = AtlasGlossaryTerm()
114+
term.relationship_status = "DELETED"
115+
table.attributes.meanings = [term]
116+
mock_methods["get_asset_by_guid"].return_value = table
117+
mock_methods["upsert"].return_value.assets_updated.return_value = [table]
118+
client = AtlanClient()
119+
guid = "123"
120+
terms = [AtlasGlossaryTerm()]
121+
122+
assert (
123+
asset := client.append_terms(guid=guid, asset_type=asset_type, terms=terms)
124+
)
125+
assert asset.terms == terms
126+
127+
128+
def test_append_with_valid_guid_when_terms_present_returns_asset_with_combined_terms(
129+
monkeypatch,
130+
):
131+
monkeypatch.setenv("ATLAN_BASE_URL", "https://name.atlan.com")
132+
monkeypatch.setenv("ATLAN_API_KEY", "abkj")
133+
asset_type = Table
134+
with patch.multiple(
135+
AtlanClient, get_asset_by_guid=DEFAULT, upsert=DEFAULT
136+
) as mock_methods:
137+
table = Table(attributes=Table.Attributes())
138+
exisiting_term = AtlasGlossaryTerm()
139+
table.attributes.meanings = [exisiting_term]
140+
mock_methods["get_asset_by_guid"].return_value = table
141+
mock_methods["upsert"].return_value.assets_updated.return_value = [table]
142+
client = AtlanClient()
143+
guid = "123"
144+
145+
new_term = AtlasGlossaryTerm()
146+
terms = [new_term]
147+
148+
assert (
149+
asset := client.append_terms(guid=guid, asset_type=asset_type, terms=terms)
150+
)
151+
assert (updated_terms := asset.terms)
152+
assert len(updated_terms) == 2
153+
assert exisiting_term in updated_terms
154+
assert new_term in updated_terms
155+
156+
157+
@pytest.mark.parametrize(
158+
"guid, qualified_name, asset_type, terms, message",
159+
[
160+
(
161+
None,
162+
None,
163+
Table,
164+
[AtlasGlossaryTerm()],
165+
"Either guid or qualified name must be specified",
166+
),
167+
(
168+
"123",
169+
None,
170+
None,
171+
[AtlasGlossaryTerm()],
172+
"1 validation error for ReplaceTerms\\nasset_type\\n none is not an allowed value ",
173+
),
174+
(
175+
"123",
176+
"default/abc",
177+
Table,
178+
[AtlasGlossaryTerm()],
179+
"Either guid or qualified_name can be be specified not both",
180+
),
181+
(
182+
"123",
183+
None,
184+
Table,
185+
None,
186+
"1 validation error for ReplaceTerms\\nterms\\n none is not an allowed value ",
187+
),
188+
],
189+
)
190+
def test_replace_terms_with_invalid_parameter_raises_valueerror(
191+
guid,
192+
qualified_name,
193+
asset_type,
194+
terms,
195+
message,
196+
monkeypatch,
197+
):
198+
monkeypatch.setenv("ATLAN_BASE_URL", "https://name.atlan.com")
199+
monkeypatch.setenv("ATLAN_API_KEY", "abkj")
200+
client = AtlanClient()
201+
202+
with pytest.raises(ValueError, match=message):
203+
client.replace_terms(
204+
guid=guid, qualified_name=qualified_name, asset_type=asset_type, terms=terms
205+
)
206+
207+
208+
def test_replace_terms(
209+
monkeypatch,
210+
):
211+
monkeypatch.setenv("ATLAN_BASE_URL", "https://name.atlan.com")
212+
monkeypatch.setenv("ATLAN_API_KEY", "abkj")
213+
asset_type = Table
214+
with patch.multiple(
215+
AtlanClient, get_asset_by_guid=DEFAULT, upsert=DEFAULT
216+
) as mock_methods:
217+
table = Table()
218+
mock_methods["get_asset_by_guid"].return_value = table
219+
mock_methods["upsert"].return_value.assets_updated.return_value = [table]
220+
client = AtlanClient()
221+
guid = "123"
222+
terms = [AtlasGlossaryTerm()]
223+
224+
assert (
225+
asset := client.replace_terms(guid=guid, asset_type=asset_type, terms=terms)
226+
)
227+
assert asset.terms == terms
228+
229+
230+
@pytest.mark.parametrize(
231+
"guid, qualified_name, asset_type, terms, message",
232+
[
233+
(
234+
None,
235+
None,
236+
Table,
237+
[AtlasGlossaryTerm()],
238+
"Either guid or qualified name must be specified",
239+
),
240+
(
241+
"123",
242+
None,
243+
None,
244+
[AtlasGlossaryTerm()],
245+
"1 validation error for RemoveTerms\\nasset_type\\n none is not an allowed value ",
246+
),
247+
(
248+
"123",
249+
"default/abc",
250+
Table,
251+
[AtlasGlossaryTerm()],
252+
"Either guid or qualified_name can be be specified not both",
253+
),
254+
(
255+
"123",
256+
None,
257+
Table,
258+
None,
259+
"1 validation error for RemoveTerms\\nterms\\n none is not an allowed value ",
260+
),
261+
(
262+
"123",
263+
None,
264+
Table,
265+
[],
266+
"A list of terms to remove must be specified",
267+
),
268+
],
269+
)
270+
def test_remove_terms_with_invalid_parameter_raises_valueerror(
271+
guid,
272+
qualified_name,
273+
asset_type,
274+
terms,
275+
message,
276+
monkeypatch,
277+
):
278+
monkeypatch.setenv("ATLAN_BASE_URL", "https://name.atlan.com")
279+
monkeypatch.setenv("ATLAN_API_KEY", "abkj")
280+
client = AtlanClient()
281+
282+
with pytest.raises(ValueError, match=message):
283+
client.remove_terms(
284+
guid=guid, qualified_name=qualified_name, asset_type=asset_type, terms=terms
285+
)
286+
287+
288+
def test_remove_with_valid_guid_when_terms_present_returns_asset_with_terms_removed(
289+
monkeypatch,
290+
):
291+
monkeypatch.setenv("ATLAN_BASE_URL", "https://name.atlan.com")
292+
monkeypatch.setenv("ATLAN_API_KEY", "abkj")
293+
asset_type = Table
294+
with patch.multiple(
295+
AtlanClient, get_asset_by_guid=DEFAULT, upsert=DEFAULT
296+
) as mock_methods:
297+
table = Table(attributes=Table.Attributes())
298+
exisiting_term = AtlasGlossaryTerm()
299+
exisiting_term.guid = "b4113341-251b-4adc-81fb-2420501c30e6"
300+
other_term = AtlasGlossaryTerm()
301+
other_term.guid = "b267858d-8316-4c41-a56a-6e9b840cef4a"
302+
table.attributes.meanings = [exisiting_term, other_term]
303+
mock_methods["get_asset_by_guid"].return_value = table
304+
mock_methods["upsert"].return_value.assets_updated.return_value = [table]
305+
client = AtlanClient()
306+
guid = "123"
307+
308+
assert (
309+
asset := client.remove_terms(
310+
guid=guid, asset_type=asset_type, terms=[exisiting_term]
311+
)
312+
)
313+
assert (updated_terms := asset.terms)
314+
assert len(updated_terms) == 1
315+
assert other_term in updated_terms

‎tests/unit/test_model.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,7 @@
132132
"Optional[QuickSightAnalysisStatus]": QuickSightAnalysisStatus.CREATION_FAILED,
133133
"Optional[QuickSightDatasetImportMode]": QuickSightDatasetImportMode.SPICE,
134134
"Optional[list[KafkaTopicConsumption]]": [KafkaTopicConsumption()],
135+
"list[AtlasGlossaryTerm]": [AtlasGlossaryTerm()],
135136
}
136137

137138

@@ -1321,7 +1322,7 @@ def test_attributes(clazz, property_name, attribute_value):
13211322
)
13221323
assert attribute_value == local_ns["ret_value"]
13231324
exec(
1324-
f"ret_value = sut.attributes.{property_name}",
1325+
f"ret_value = sut.attributes.{property_name if property_name != 'terms' else 'meanings'}",
13251326
{"sut": sut, "property_name": property_name},
13261327
local_ns,
13271328
)

0 commit comments

Comments
 (0)
Please sign in to comment.