def test_handle_nexttoken_and_retries():
def mock_function_set_next_token(**kwargs):
if not kwargs.get('NextToken'):
return {'mock-key': ['1st-pass'], 'NextToken': 'mock-token-str'}
else:
return {'mock-key': ['2nd-pass']}
def mock_function_raise_client_error(error_code):
raise ClientError(
{
'Error': {
'Code': error_code
},
},
'mock_function',
)
def mock_function_raise_value_error():
raise ValueError('this is a value error')
org = orgs.Org(MASTER_ACCOUNT_ID, ORG_ACCESS_ROLE)
collector = utils.handle_nexttoken_and_retries(
obj=org,
collector_key='mock-key',
function=mock_function_set_next_token,
kwargs=dict(),
)
assert collector == ['1st-pass', '2nd-pass']
exception_name = 'TooManyRequestsException'
with pytest.raises(ClientError) as e:
collector = utils.handle_nexttoken_and_retries(
obj=org,
collector_key='mock-key',
function=mock_function_raise_client_error,
kwargs=dict(error_code=exception_name),
)
assert e.value.response['Error']['Code'] == exception_name
exception_name = 'SomeOtherException'
with pytest.raises(ClientError) as e:
collector = utils.handle_nexttoken_and_retries(
obj=org,
collector_key='mock-key',
function=mock_function_raise_client_error,
kwargs=dict(error_code=exception_name),
)
assert e.value.response['Error']['Code'] == exception_name
with pytest.raises(ValueError) as e:
collector = utils.handle_nexttoken_and_retries(
obj=org,
collector_key='mock-key',
function=mock_function_raise_value_error,
kwargs=dict(),
)
def load_targets(self):
message = {
'FILE': __file__.split('/')[-1],
'CLASS': self.__class__.__name__,
'METHOD': inspect.stack()[0][3],
'policy': self.name,
}
self.logger.info(message)
self.targets = utils.handle_nexttoken_and_retries(
obj=self,
collector_key='Targets',
function=self.client.list_targets_for_policy,
kwargs=dict(PolicyId=self.id, ))
def _load_accounts(self):
message = {
'FILE': __file__.split('/')[-1],
'CLASS': self.__class__.__name__,
'METHOD': inspect.stack()[0][3],
}
self.logger.info(message)
accounts = utils.handle_nexttoken_and_retries(
obj=self,
collector_key='Accounts',
function=self.client.list_accounts,
)
# skip accounts with no 'Name' key, as these are not fully created yet.
accounts = [account for account in accounts if 'Name' in account]
def make_org_account_object(account, org):
message = {
'FILE': __file__.split('/')[-1],
'CLASS': self.__class__.__name__,
'METHOD': inspect.stack()[0][3],
'account_id': account['Id'],
'account_name': account['Name'],
}
self.logger.info(message)
try:
org_account = OrgAccount(
org,
name=account['Name'],
id=account['Id'],
email=account['Email'],
status=account['Status'],
)
org_account.get_parent_id()
org_account.load_attached_policy_ids()
org.accounts.append(org_account)
except Exception: # pragma: no cover
org._exc_info = sys.exc_info()
utils.queue_threads(
accounts,
make_org_account_object,
func_args=(self,),
logger=self.logger,
)
if self._exc_info: # pragma: no cover
raise self._exc_info[1].with_traceback(self._exc_info[2])
def load_attached_policy_ids(self):
message = {
'FILE': __file__.split('/')[-1],
'CLASS': self.__class__.__name__,
'METHOD': inspect.stack()[0][3],
'object_id': self.id,
'object_name': self.name,
}
self.logger.info(message)
policies = utils.handle_nexttoken_and_retries(
obj=self,
collector_key='Policies',
function=self.client.list_policies_for_target,
kwargs=dict(
TargetId=self.id,
Filter='SERVICE_CONTROL_POLICY',
))
self.attached_policy_ids = [p['Id'] for p in policies]
def get_parent_id(self):
'''
Set the parent id in OrgObject. Thread freindly.
'''
message = {
'FILE': __file__.split('/')[-1],
'CLASS': self.__class__.__name__,
'METHOD': inspect.stack()[0][3],
'object_id': self.id,
'object_name': self.name,
}
self.logger.info(message)
parents = utils.handle_nexttoken_and_retries(
obj=self,
collector_key='Parents',
function=self.client.list_parents,
kwargs=dict(ChildId=self.id, ))
self.parent_id = parents[0]['Id']
def _load_policies(self):
message = {
'FILE': __file__.split('/')[-1],
'CLASS': self.__class__.__name__,
'METHOD': inspect.stack()[0][3],
}
self.logger.info(message)
policies = utils.handle_nexttoken_and_retries(
obj=self,
collector_key='Policies',
function=self.client.list_policies,
kwargs=dict(Filter='SERVICE_CONTROL_POLICY'),
)
def make_org_policy_object(policy, org):
message = {
'FILE': __file__.split('/')[-1],
'CLASS': self.__class__.__name__,
'METHOD': inspect.stack()[0][3],
'policy': policy,
}
self.logger.info(message)
try:
org_policy = OrgPolicy(
org,
name=policy['Name'],
id=policy['Id'],
)
org_policy.load_targets()
org.policies.append(org_policy)
except Exception: # pragma: no cover
org._exc_info = sys.exc_info()
utils.queue_threads(
policies,
make_org_policy_object,
func_args=(self,),
logger=self.logger,
)
if self._exc_info: # pragma: no cover
raise self._exc_info[1].with_traceback(self._exc_info[2])
def _recurse_organization(self, parent_id):
message = {
'FILE': __file__.split('/')[-1],
'CLASS': self.__class__.__name__,
'METHOD': inspect.stack()[0][3],
}
self.logger.info(message)
org_units = utils.handle_nexttoken_and_retries(
obj=self,
collector_key='OrganizationalUnits',
function=self.client.list_organizational_units_for_parent,
kwargs=dict(ParentId=parent_id),
)
for ou in org_units:
org_unit = OrganizationalUnit(
self,
name=ou['Name'],
id=ou['Id'],
parent_id=parent_id,
)
org_unit.load_attached_policy_ids()
self.org_units.append(org_unit)
self._recurse_organization(ou['Id'])