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'])