11# SPDX-License-Identifier: Apache-2.0
22# Copyright 2022 Atlan Pte. Ltd.
3- from typing import Generator
3+ from typing import Generator , Optional
44
55import pytest
6+ from pydantic .v1 import StrictStr
67
78from pyatlan .client .atlan import AtlanClient
8- from pyatlan .model .assets import AtlasGlossary , AuthPolicy , Connection , Persona
9+ from pyatlan .model .assets import (
10+ AccessControl ,
11+ AtlasGlossary ,
12+ AuthPolicy ,
13+ Connection ,
14+ Persona ,
15+ )
16+ from pyatlan .model .core import AtlanObject
917from pyatlan .model .enums import (
1018 AssetSidebarTab ,
1119 AtlanConnectorType ,
1422 PersonaGlossaryAction ,
1523 PersonaMetadataAction ,
1624)
25+ from pyatlan .model .fluent_search import CompoundQuery , FluentSearch
1726from tests .integration .client import TestId , delete_asset
1827from tests .integration .connection_test import create_connection
1928from tests .integration .glossary_test import create_glossary
@@ -56,6 +65,16 @@ def persona(
5665 delete_asset (client , guid = p .guid , asset_type = Persona )
5766
5867
68+ class PolicyInfo (AtlanObject ):
69+ guid : Optional [str ]
70+ name : Optional [str ]
71+
72+
73+ @pytest .fixture (scope = "module" )
74+ def policy_info () -> PolicyInfo :
75+ return PolicyInfo (guid = None , name = None )
76+
77+
5978def test_persona (
6079 client : AtlanClient ,
6180 persona : Persona ,
@@ -162,6 +181,7 @@ def test_retrieve_persona(
162181 persona : Persona ,
163182 connection : Connection ,
164183 glossary : AtlasGlossary ,
184+ policy_info : PolicyInfo ,
165185):
166186 assert persona .qualified_name
167187 one = client .asset .get_by_qualified_name (
@@ -187,6 +207,9 @@ def test_retrieve_persona(
187207 full = client .asset .get_by_guid (
188208 guid = policy .guid , asset_type = AuthPolicy , ignore_relationships = False
189209 )
210+ if policy_info .guid is None and policy_info .name is None :
211+ policy_info .guid = full .guid
212+ policy_info .name = full .name
190213 assert full
191214 sub_cat = full .policy_sub_category
192215 assert sub_cat
@@ -211,3 +234,45 @@ def test_retrieve_persona(
211234 assert PersonaGlossaryAction .UPDATE in full .policy_actions
212235 assert full .policy_resources
213236 assert f"entity:{ glossary .qualified_name } " in full .policy_resources
237+
238+
239+ @pytest .mark .order (after = "test_retrieve_persona" )
240+ def test_update_policy (
241+ client : AtlanClient ,
242+ policy_info : PolicyInfo ,
243+ ):
244+ assert policy_info .guid
245+ assert policy_info .name
246+ request = (
247+ FluentSearch ()
248+ .where (FluentSearch .asset_type (AuthPolicy ))
249+ .where (AuthPolicy .POLICY_CATEGORY .eq (StrictStr ("persona" )))
250+ .where (AuthPolicy .NAME .eq (policy_info .name ))
251+ .where (AuthPolicy .GUID .eq (policy_info .guid ))
252+ .where (CompoundQuery .active_assets ())
253+ .include_on_results (AuthPolicy .POLICY_CATEGORY )
254+ .include_on_results (AuthPolicy .NAME )
255+ .include_on_results (AuthPolicy .POLICY_SERVICE_NAME )
256+ .include_on_results (AuthPolicy .ACCESS_CONTROL )
257+ .include_on_results (AuthPolicy .POLICY_ACTIONS )
258+ .include_on_results (AuthPolicy .POLICY_RESOURCES )
259+ .include_on_results (AuthPolicy .CONNECTION_QUALIFIED_NAME )
260+ .include_on_results (AuthPolicy .POLICY_TYPE )
261+ .include_on_results (AuthPolicy .POLICY_SUB_CATEGORY )
262+ .include_on_relations (AccessControl .IS_ACCESS_CONTROL_ENABLED )
263+ .include_on_relations (AccessControl .NAME )
264+ ).to_request ()
265+ to_update = client .asset .search (request )
266+
267+ assert to_update .count == 1
268+ policy = to_update .current_page ()[0 ]
269+ assert policy
270+ policy .name = f"Updated policy ({ MODULE_NAME } )"
271+
272+ response = client .asset .save (policy )
273+ assert response
274+ updated = response .assets_updated (asset_type = AuthPolicy )
275+ assert updated
276+ assert len (updated ) == 1
277+ assert updated [0 ].guid == policy .guid
278+ assert updated [0 ].name == f"Updated policy ({ MODULE_NAME } )"
0 commit comments