Skip to content

Commit 6254025

Browse files
Aryamanz29Aryamanz29
committed
[change] Updated integration tests as per refactoring changes
1 parent 2407819 commit 6254025

9 files changed

Lines changed: 111 additions & 57 deletions

‎tests/integration/connection_test.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,10 @@ def create_connection(
1717
) -> Connection:
1818
admin_role_guid = str(client.role_cache.get_id_for_name("$admin"))
1919
to_create = Connection.create(
20-
name=name, connector_type=connector_type, admin_roles=[admin_role_guid]
20+
client=client,
21+
name=name,
22+
connector_type=connector_type,
23+
admin_roles=[admin_role_guid],
2124
)
2225
response = client.asset.save(to_create)
2326
result = response.assets_created(asset_type=Connection)[0]

‎tests/integration/custom_metadata_test.py‎

Lines changed: 67 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,7 @@ def update_enum(
118118
client: AtlanClient, name: str, values: List[str], replace_existing: bool = False
119119
) -> EnumDef:
120120
enum_def = EnumDef.update(
121-
name=name, values=values, replace_existing=replace_existing
121+
client=client, name=name, values=values, replace_existing=replace_existing
122122
)
123123
r = client.typedef.update(enum_def)
124124
return r.enum_defs[0]
@@ -143,24 +143,29 @@ def cm_ipr(
143143
) -> Generator[CustomMetadataDef, None, None]:
144144
attribute_defs = [
145145
AttributeDef.create(
146+
client=client,
146147
display_name=CM_ATTR_IPR_LICENSE,
147148
attribute_type=AtlanCustomAttributePrimitiveType.STRING,
148149
description=ATTRIBUTE_DESCRIPTION,
149150
**limit_attribute_applicability_kwargs,
150151
),
151152
AttributeDef.create(
153+
client=client,
152154
display_name=CM_ATTR_IPR_VERSION,
153155
attribute_type=AtlanCustomAttributePrimitiveType.DECIMAL,
154156
),
155157
AttributeDef.create(
158+
client=client,
156159
display_name=CM_ATTR_IPR_MANDATORY,
157160
attribute_type=AtlanCustomAttributePrimitiveType.BOOLEAN,
158161
),
159162
AttributeDef.create(
163+
client=client,
160164
display_name=CM_ATTR_IPR_DATE,
161165
attribute_type=AtlanCustomAttributePrimitiveType.DATE,
162166
),
163167
AttributeDef.create(
168+
client=client,
164169
display_name=CM_ATTR_IPR_URL,
165170
attribute_type=AtlanCustomAttributePrimitiveType.URL,
166171
),
@@ -230,28 +235,33 @@ def cm_raci(
230235
client: AtlanClient,
231236
) -> Generator[CustomMetadataDef, None, None]:
232237
TEST_MULTI_VALUE_USING_SETTER = AttributeDef.create(
238+
client=client,
233239
display_name=CM_ATTR_RACI_INFORMED,
234240
attribute_type=AtlanCustomAttributePrimitiveType.GROUPS,
235241
)
236242
assert TEST_MULTI_VALUE_USING_SETTER and TEST_MULTI_VALUE_USING_SETTER.options
237243
TEST_MULTI_VALUE_USING_SETTER.options.multi_value_select = True
238244
attribute_defs = [
239245
AttributeDef.create(
246+
client=client,
240247
display_name=CM_ATTR_RACI_RESPONSIBLE,
241248
attribute_type=AtlanCustomAttributePrimitiveType.USERS,
242249
multi_valued=True,
243250
),
244251
AttributeDef.create(
252+
client=client,
245253
display_name=CM_ATTR_RACI_ACCOUNTABLE,
246254
attribute_type=AtlanCustomAttributePrimitiveType.USERS,
247255
),
248256
AttributeDef.create(
257+
client=client,
249258
display_name=CM_ATTR_RACI_CONSULTED,
250259
attribute_type=AtlanCustomAttributePrimitiveType.GROUPS,
251260
multi_valued=True,
252261
),
253262
TEST_MULTI_VALUE_USING_SETTER,
254263
AttributeDef.create(
264+
client=client,
255265
display_name=CM_ATTR_RACI_EXTRA,
256266
attribute_type=AtlanCustomAttributePrimitiveType.STRING,
257267
),
@@ -394,14 +404,17 @@ def cm_dq(
394404
) -> Generator[CustomMetadataDef, None, None]:
395405
attribute_defs = [
396406
AttributeDef.create(
407+
client=client,
397408
display_name=CM_ATTR_QUALITY_COUNT,
398409
attribute_type=AtlanCustomAttributePrimitiveType.INTEGER,
399410
),
400411
AttributeDef.create(
412+
client=client,
401413
display_name=CM_ATTR_QUALITY_SQL,
402414
attribute_type=AtlanCustomAttributePrimitiveType.SQL,
403415
),
404416
AttributeDef.create(
417+
client=client,
405418
display_name=CM_ATTR_QUALITY_TYPE,
406419
attribute_type=AtlanCustomAttributePrimitiveType.OPTIONS,
407420
options_name=DQ_ENUM,
@@ -518,7 +531,7 @@ def test_add_term_cm_raci(
518531
groups: List[AtlanGroup],
519532
):
520533
cm_name = CM_RACI
521-
raci_attrs = CustomMetadataDict(cm_name)
534+
raci_attrs = CustomMetadataDict(client=client, name=cm_name)
522535
_validate_raci_empty(raci_attrs)
523536
group1, group2 = _get_groups(client)
524537
raci_attrs[CM_ATTR_RACI_RESPONSIBLE] = [FIXED_USER]
@@ -528,7 +541,9 @@ def test_add_term_cm_raci(
528541
client.asset.update_custom_metadata_attributes(term.guid, raci_attrs)
529542
t = client.asset.retrieve_minimal(guid=term.guid, asset_type=AtlasGlossaryTerm)
530543
assert t
531-
_validate_raci_attributes(client, t.get_custom_metadata(name=cm_name))
544+
_validate_raci_attributes(
545+
client, t.get_custom_metadata(client=client, name=cm_name)
546+
)
532547

533548

534549
def test_add_term_cm_ipr(
@@ -537,7 +552,7 @@ def test_add_term_cm_ipr(
537552
term: AtlasGlossaryTerm,
538553
):
539554
cm_name = CM_IPR
540-
ipr_attrs = CustomMetadataDict(cm_name)
555+
ipr_attrs = CustomMetadataDict(client=client, name=cm_name)
541556
_validate_ipr_empty(ipr_attrs)
542557
ipr_attrs[CM_ATTR_IPR_LICENSE] = "CC BY"
543558
ipr_attrs[CM_ATTR_IPR_VERSION] = 2.0
@@ -548,7 +563,7 @@ def test_add_term_cm_ipr(
548563
client.asset.update_custom_metadata_attributes(term.guid, ipr_attrs)
549564
t = client.asset.retrieve_minimal(guid=term.guid, asset_type=AtlasGlossaryTerm)
550565
assert t
551-
_validate_ipr_attributes(t.get_custom_metadata(name=cm_name))
566+
_validate_ipr_attributes(t.get_custom_metadata(client=client, name=cm_name))
552567

553568

554569
def test_add_term_cm_dq(
@@ -557,15 +572,15 @@ def test_add_term_cm_dq(
557572
term: AtlasGlossaryTerm,
558573
):
559574
cm_name = CM_QUALITY
560-
dq_attrs = CustomMetadataDict(cm_name)
575+
dq_attrs = CustomMetadataDict(client=client, name=cm_name)
561576
_validate_dq_empty(dq_attrs)
562577
dq_attrs[CM_ATTR_QUALITY_COUNT] = 42
563578
dq_attrs[CM_ATTR_QUALITY_SQL] = "SELECT * from SOMEWHERE;"
564579
dq_attrs[CM_ATTR_QUALITY_TYPE] = "Completeness"
565580
client.asset.update_custom_metadata_attributes(term.guid, dq_attrs)
566581
t = client.asset.retrieve_minimal(guid=term.guid, asset_type=AtlasGlossaryTerm)
567582
assert t
568-
_validate_dq_attributes(t.get_custom_metadata(name=cm_name))
583+
_validate_dq_attributes(t.get_custom_metadata(client=client, name=cm_name))
569584

570585

571586
@pytest.mark.order(after="test_add_term_cm_dq")
@@ -575,15 +590,19 @@ def test_update_term_cm_ipr(
575590
term: AtlasGlossaryTerm,
576591
):
577592
cm_name = CM_IPR
578-
ipr = CustomMetadataDict(cm_name)
593+
ipr = CustomMetadataDict(client=client, name=cm_name)
579594
# Note: MUST access the getter / setter, not the underlying store
580595
ipr[CM_ATTR_IPR_MANDATORY] = False
581596
client.asset.update_custom_metadata_attributes(term.guid, ipr)
582597
t = client.asset.retrieve_minimal(guid=term.guid, asset_type=AtlasGlossaryTerm)
583598
assert t
584-
_validate_ipr_attributes(t.get_custom_metadata(name=cm_name), mandatory=False)
585-
_validate_raci_attributes(client, t.get_custom_metadata(name=CM_RACI))
586-
_validate_dq_attributes(t.get_custom_metadata(name=CM_QUALITY))
599+
_validate_ipr_attributes(
600+
t.get_custom_metadata(client=client, name=cm_name), mandatory=False
601+
)
602+
_validate_raci_attributes(
603+
client, t.get_custom_metadata(client=client, name=CM_RACI)
604+
)
605+
_validate_dq_attributes(t.get_custom_metadata(client=client, name=CM_QUALITY))
587606

588607

589608
@pytest.mark.order(after="test_update_term_cm_ipr")
@@ -592,7 +611,7 @@ def test_replace_term_cm_raci(
592611
cm_raci: CustomMetadataDef,
593612
term: AtlasGlossaryTerm,
594613
):
595-
raci = CustomMetadataDict(CM_RACI)
614+
raci = CustomMetadataDict(client=client, name=CM_RACI)
596615
group1, group2 = _get_groups(client)
597616
raci[CM_ATTR_RACI_RESPONSIBLE] = [FIXED_USER]
598617
raci[CM_ATTR_RACI_ACCOUNTABLE] = FIXED_USER
@@ -601,9 +620,13 @@ def test_replace_term_cm_raci(
601620
client.asset.replace_custom_metadata(term.guid, raci)
602621
t = client.asset.retrieve_minimal(guid=term.guid, asset_type=AtlasGlossaryTerm)
603622
assert t
604-
_validate_raci_attributes_replacement(client, t.get_custom_metadata(name=CM_RACI))
605-
_validate_ipr_attributes(t.get_custom_metadata(name=CM_IPR), mandatory=False)
606-
_validate_dq_attributes(t.get_custom_metadata(name=CM_QUALITY))
623+
_validate_raci_attributes_replacement(
624+
client, t.get_custom_metadata(client=client, name=CM_RACI)
625+
)
626+
_validate_ipr_attributes(
627+
t.get_custom_metadata(client=client, name=CM_IPR), mandatory=False
628+
)
629+
_validate_dq_attributes(t.get_custom_metadata(client=client, name=CM_QUALITY))
607630

608631

609632
@pytest.mark.order(after="test_replace_term_cm_raci")
@@ -612,13 +635,15 @@ def test_replace_term_cm_ipr(
612635
cm_ipr: CustomMetadataDef,
613636
term: AtlasGlossaryTerm,
614637
):
615-
term_cm_ipr = CustomMetadataDict(CM_IPR)
638+
term_cm_ipr = CustomMetadataDict(client=client, name=CM_IPR)
616639
client.asset.replace_custom_metadata(term.guid, term_cm_ipr)
617640
t = client.asset.retrieve_minimal(guid=term.guid, asset_type=AtlasGlossaryTerm)
618641
assert t
619-
_validate_raci_attributes_replacement(client, t.get_custom_metadata(name=CM_RACI))
620-
_validate_dq_attributes(t.get_custom_metadata(name=CM_QUALITY))
621-
_validate_ipr_empty(t.get_custom_metadata(name=CM_IPR))
642+
_validate_raci_attributes_replacement(
643+
client, t.get_custom_metadata(client=client, name=CM_RACI)
644+
)
645+
_validate_dq_attributes(t.get_custom_metadata(client=client, name=CM_QUALITY))
646+
_validate_ipr_empty(t.get_custom_metadata(client=client, name=CM_IPR))
622647

623648

624649
@pytest.mark.order(after="test_replace_term_cm_ipr")
@@ -638,7 +663,11 @@ def test_search_by_any_accountable(
638663
FluentSearch(_includes_on_results=attributes)
639664
.where(CompoundQuery.active_assets())
640665
.where(CompoundQuery.asset_type(AtlasGlossaryTerm))
641-
.where(CustomMetadataField(CM_RACI, CM_ATTR_RACI_ACCOUNTABLE).has_any_value())
666+
.where(
667+
CustomMetadataField(
668+
client, CM_RACI, CM_ATTR_RACI_ACCOUNTABLE
669+
).has_any_value()
670+
)
642671
.include_on_relations(Asset.NAME)
643672
).to_request()
644673
response = client.asset.search(criteria=request)
@@ -658,7 +687,7 @@ def test_search_by_any_accountable(
658687
assert anchor
659688
assert anchor.name == glossary.name
660689
_validate_raci_attributes_replacement(
661-
client, t.get_custom_metadata(name=CM_RACI)
690+
client, t.get_custom_metadata(client=client, name=CM_RACI)
662691
)
663692

664693

@@ -673,7 +702,11 @@ def test_search_by_specific_accountable(
673702
FluentSearch()
674703
.where(CompoundQuery.active_assets())
675704
.where(CompoundQuery.asset_type(AtlasGlossaryTerm))
676-
.where(CustomMetadataField(CM_RACI, CM_ATTR_RACI_ACCOUNTABLE).eq(FIXED_USER))
705+
.where(
706+
CustomMetadataField(client, CM_RACI, CM_ATTR_RACI_ACCOUNTABLE).eq(
707+
FIXED_USER
708+
)
709+
)
677710
.include_on_results(Asset.NAME)
678711
.include_on_results(AtlasGlossaryTerm.ANCHOR)
679712
.include_on_relations(Asset.NAME)
@@ -707,9 +740,9 @@ def test_remove_term_cm_raci(
707740
client.asset.remove_custom_metadata(term.guid, cm_name=CM_RACI)
708741
t = client.asset.retrieve_minimal(guid=term.guid, asset_type=AtlasGlossaryTerm)
709742
assert t
710-
_validate_dq_attributes(t.get_custom_metadata(name=CM_QUALITY))
711-
_validate_ipr_empty(t.get_custom_metadata(name=CM_IPR))
712-
_validate_raci_empty(t.get_custom_metadata(name=CM_RACI))
743+
_validate_dq_attributes(t.get_custom_metadata(client=client, name=CM_QUALITY))
744+
_validate_ipr_empty(t.get_custom_metadata(client=client, name=CM_IPR))
745+
_validate_raci_empty(t.get_custom_metadata(client=client, name=CM_RACI))
713746

714747

715748
@pytest.mark.order(after="test_remove_term_cm_raci")
@@ -721,9 +754,9 @@ def test_remove_term_cm_ipr(
721754
client.asset.remove_custom_metadata(term.guid, cm_name=CM_IPR)
722755
t = client.asset.retrieve_minimal(guid=term.guid, asset_type=AtlasGlossaryTerm)
723756
assert t
724-
_validate_dq_attributes(t.get_custom_metadata(name=CM_QUALITY))
725-
_validate_ipr_empty(t.get_custom_metadata(name=CM_IPR))
726-
_validate_raci_empty(t.get_custom_metadata(name=CM_RACI))
757+
_validate_dq_attributes(t.get_custom_metadata(client=client, name=CM_QUALITY))
758+
_validate_ipr_empty(t.get_custom_metadata(client=client, name=CM_IPR))
759+
_validate_raci_empty(t.get_custom_metadata(client=client, name=CM_RACI))
727760

728761

729762
@pytest.mark.order(after="test_remove_term_cm_raci")
@@ -798,6 +831,7 @@ def test_recreate_attribute(client: AtlanClient, cm_raci: CustomMetadataDef):
798831
existing_attr.is_new = None
799832
updated_attrs.append(existing_attr)
800833
new_attr = AttributeDef.create(
834+
client=client,
801835
display_name=CM_ATTR_RACI_EXTRA,
802836
attribute_type=AtlanCustomAttributePrimitiveType.STRING,
803837
)
@@ -870,7 +904,7 @@ def test_update_replacing_cm(
870904
cm_dq: CustomMetadataDef,
871905
client: AtlanClient,
872906
):
873-
raci = CustomMetadataDict(CM_RACI)
907+
raci = CustomMetadataDict(client=client, name=CM_RACI)
874908
group1, group2 = _get_groups(client)
875909
raci[CM_ATTR_RACI_RESPONSIBLE] = [FIXED_USER]
876910
raci[CM_ATTR_RACI_ACCOUNTABLE] = FIXED_USER
@@ -882,7 +916,7 @@ def test_update_replacing_cm(
882916
to_update = AtlasGlossaryTerm.create_for_modification(
883917
qualified_name=term.qualified_name, name=term.name, glossary_guid=glossary.guid
884918
)
885-
to_update.set_custom_metadata(custom_metadata=raci)
919+
to_update.set_custom_metadata(custom_metadata=raci, client=client)
886920
response = client.asset.update_replacing_cm(to_update, replace_atlan_tags=False)
887921
assert response
888922
assert len(response.assets_deleted(asset_type=AtlasGlossaryTerm)) == 0
@@ -901,11 +935,11 @@ def test_update_replacing_cm(
901935
assert x
902936
assert not x.is_incomplete
903937
assert x.qualified_name == term.qualified_name
904-
raci = x.get_custom_metadata(CM_RACI)
938+
raci = x.get_custom_metadata(client=client, name=CM_RACI)
905939
_validate_raci_attributes(client, raci)
906940
assert raci[CM_ATTR_RACI_EXTRA] == "something extra..."
907-
_validate_ipr_empty(x.get_custom_metadata(CM_IPR))
908-
_validate_dq_empty(x.get_custom_metadata(CM_QUALITY))
941+
_validate_ipr_empty(x.get_custom_metadata(client=client, name=CM_IPR))
942+
_validate_dq_empty(x.get_custom_metadata(client=client, name=CM_QUALITY))
909943

910944

911945
# TODO: test entity audit retrieval and parsing, once available

‎tests/integration/data_mesh_test.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -222,6 +222,7 @@ def data_domain_cm(
222222
assert domain.qualified_name
223223
attribute_defs = [
224224
AttributeDef.create(
225+
client=client,
225226
display_name=DD_ATTR,
226227
attribute_type=AtlanCustomAttributePrimitiveType.STRING,
227228
applicable_domain_types={"DataDomain", "DataProduct"},
@@ -491,7 +492,7 @@ def test_product_get_assets(client: AtlanClient, product: DataProduct):
491492
)
492493
assert test_product
493494
assert test_product.data_product_assets_d_s_l
494-
asset_list = test_product.get_assets()
495+
asset_list = test_product.get_assets(client=client)
495496
assert asset_list.count and asset_list.count > 0
496497
TOTAL_ASSETS = asset_list.count
497498
counter = 0

‎tests/integration/purpose_test.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,13 +196,15 @@ def test_add_policies_to_purpose(
196196
token: ApiToken,
197197
):
198198
metadata = Purpose.create_metadata_policy(
199+
client=client,
199200
name="Simple read access",
200201
purpose_id=purpose.guid,
201202
policy_type=AuthPolicyType.ALLOW,
202203
actions={PurposeMetadataAction.READ},
203204
all_users=True,
204205
)
205206
data = Purpose.create_data_policy(
207+
client=client,
206208
name="Mask the data",
207209
purpose_id=purpose.guid,
208210
policy_type=AuthPolicyType.DATA_MASK,

0 commit comments

Comments
 (0)