def check_limits(self, verbose=False): """ Run the actual usage and limit check, with overrides. see: http://awslimitchecker.readthedocs.org/en/latest/python_usage.html#ci-deployment-checks """ # instantiate the class checker = AwsLimitChecker() # set your overrides checker.set_threshold_overrides(AWS_THRESHOLD_OVERRIDES) checker.set_limit_overrides(AWS_LIMIT_OVERRIDES) print("Checking AWS resource usage; WARNING threshold {w}% of " "limit, CRITICAL threshold {c}% of limit".format( w=checker.warning_threshold, c=checker.critical_threshold)) # check usage against thresholds # if we didn't support verbose output, we could just iterate the return # value of this to be a bit more efficient. checker.check_thresholds() # save state for exit code and summary warnings = [] criticals = [] # iterate the results for service, svc_limits in sorted(checker.get_limits().items()): for limit_name, limit in sorted(svc_limits.items()): have_alarms = False # check warnings and criticals for each Limit for warn in limit.get_warnings(): warnings.append("{service} '{limit_name}' usage " "({u}) exceeds warning threshold " "(limit={l})".format( service=service, limit_name=limit_name, u=str(warn), l=limit.get_limit(), )) have_alarms = True for crit in limit.get_criticals(): criticals.append("{service} '{limit_name}' usage " "({u}) exceeds critical threshold" " (limit={l})".format( service=service, limit_name=limit_name, u=str(crit), l=limit.get_limit(), )) have_alarms = True if not have_alarms and verbose: print( "{service} '{limit_name}' OK: {u} (limit={l})".format( service=service, limit_name=limit_name, u=limit.get_current_usage_str(), l=limit.get_limit())) if verbose: print("\n\n") return (warnings, criticals)
def check_limits(self, acct_id, region_name, role_name=None, limit_overrides={}, threshold_overrides={}): """ Run the actual usage and limit check, with overrides, against a specific account in a specific region, optionally assuming a role in the account and optionally setting limit and/or threshold overrides. Return a 2-tuple of lists, warning strings and critical strings. see: http://awslimitchecker.readthedocs.org/en/latest/python_usage.html :returns: 2-tuple of lists of strings, warnings and criticals """ # instantiate the class if role_name is not None: checker = AwsLimitChecker( account_id=acct_id, region=region_name, account_role=role_name ) else: checker = AwsLimitChecker(region=region_name) # set your overrides if len(threshold_overrides) > 0: checker.set_threshold_overrides(threshold_overrides) if len(limit_overrides) > 0: checker.set_limit_overrides(limit_overrides) # check usage against thresholds checker.check_thresholds() # save state for exit code and summary warnings = [] criticals = [] # iterate the results for service, svc_limits in sorted(checker.get_limits().items()): for limit_name, limit in sorted(svc_limits.items()): # check warnings and criticals for each Limit for warn in limit.get_warnings(): warnings.append(colored("{service} '{limit_name}' usage " "({u}) exceeds warning threshold " "(limit={l})".format( service=service, limit_name=limit_name, u=str(warn), l=limit.get_limit(), ), 'yellow')) for crit in limit.get_criticals(): criticals.append(colored("{service} '{limit_name}' usage " "({u}) exceeds critical threshold" " (limit={l})".format( service=service, limit_name=limit_name, u=str(crit), l=limit.get_limit(), ), 'red')) return warnings, criticals
class TestAwsLimitChecker(object): def setup(self): self.mock_ver_info = Mock(release='1.2.3', url='http://myurl', commit='abcd', tag='mytag', version_str='1.2.3@mytag') self.mock_svc1 = Mock(spec_set=_AwsService) self.mock_svc2 = Mock(spec_set=ApiServiceSpec) self.mock_foo = Mock(spec_set=_AwsService) self.mock_bar = Mock(spec_set=_AwsService) self.mock_ta = Mock(spec_set=TrustedAdvisor) self.mock_foo.return_value = self.mock_svc1 self.mock_bar.return_value = self.mock_svc2 self.svcs = {'SvcFoo': self.mock_foo, 'SvcBar': self.mock_bar} with patch.dict('%s._services' % pbm, values=self.svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, _get_latest_version=DEFAULT, ServiceQuotasClient=DEFAULT, autospec=True, ) as mocks: self.mock_logger = mocks['logger'] self.mock_version = mocks['_get_version_info'] self.mock_ta_constr = mocks['TrustedAdvisor'] self.mock_glv = mocks['_get_latest_version'] self.mock_quotas = mocks['ServiceQuotasClient'] mocks['TrustedAdvisor'].return_value = self.mock_ta mocks['_get_latest_version'].return_value = None self.mock_version.return_value = self.mock_ver_info self.cls = AwsLimitChecker(check_version=False) def test_init(self): # dict should be of _AwsService instances services = {'SvcFoo': self.mock_svc1, 'SvcBar': self.mock_svc2} assert self.cls.services == services # _AwsService instances should exist, but have no other calls assert self.mock_foo.mock_calls == [ call(80, 99, {'region_name': None}, self.mock_quotas.return_value) ] assert self.mock_bar.mock_calls == [ call(80, 99, {'region_name': None}, self.mock_quotas.return_value) ] assert self.mock_ta_constr.mock_calls == [ call(services, {'region_name': None}, ta_api_region='us-east-1', ta_refresh_mode=None, ta_refresh_timeout=None) ] assert self.mock_svc1.mock_calls == [] assert self.mock_svc2.mock_calls == [] assert self.cls.ta == self.mock_ta assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info assert self.mock_glv.mock_calls == [] assert self.mock_logger.mock_calls == [ call.debug('Connecting to region %s', None) ] assert self.cls.role_partition == 'aws' assert self.mock_quotas.mock_calls == [call({'region_name': None})] def test_init_AGPL_message(self, capsys): # get rid of the class self.cls = None # clear out/err out, err = capsys.readouterr() # run setup again self.setup() # check out/err out, err = capsys.readouterr() assert out == '' assert (err) == ( "awslimitchecker 1.2.3@mytag is AGPL-licensed free software; " "all users have a right to the full source code of " "this version. See <http://myurl>\n") def test_check_version_old(self): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, _get_latest_version=DEFAULT, autospec=True, ) as mocks: mocks['_get_version_info'].return_value = self.mock_ver_info mocks['_get_latest_version'].return_value = '3.4.5' AwsLimitChecker() assert mocks['_get_latest_version'].mock_calls == [call()] assert mocks['logger'].mock_calls == [ call.warning( 'You are running awslimitchecker %s, but the latest version' ' is %s; please consider upgrading.', '1.2.3', '3.4.5'), call.debug('Connecting to region %s', None) ] def test_check_version_not_old(self): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, _get_latest_version=DEFAULT, autospec=True, ) as mocks: mocks['_get_version_info'].return_value = self.mock_ver_info mocks['_get_latest_version'].return_value = None AwsLimitChecker() assert mocks['_get_latest_version'].mock_calls == [call()] assert mocks['logger'].mock_calls == [ call.debug('Connecting to region %s', None) ] def test_init_thresholds(self): mock_svc1 = Mock(spec_set=_AwsService) mock_svc2 = Mock(spec_set=_AwsService) mock_foo = Mock(spec_set=_AwsService) mock_bar = Mock(spec_set=_AwsService) mock_ta = Mock(spec_set=TrustedAdvisor) mock_foo.return_value = mock_svc1 mock_bar.return_value = mock_svc2 svcs = {'SvcFoo': mock_foo, 'SvcBar': mock_bar} with patch.dict('%s._services' % pbm, values=svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, _get_latest_version=DEFAULT, ServiceQuotasClient=DEFAULT, autospec=True, ) as mocks: mock_version = mocks['_get_version_info'] mock_version.return_value = self.mock_ver_info mock_ta_constr = mocks['TrustedAdvisor'] mocks['TrustedAdvisor'].return_value = mock_ta mocks['_get_latest_version'].return_value = None cls = AwsLimitChecker( warning_threshold=5, critical_threshold=22, ) # dict should be of _AwsService instances services = {'SvcFoo': mock_svc1, 'SvcBar': mock_svc2} assert cls.services == services # _AwsService instances should exist, but have no other calls assert mock_foo.mock_calls == [ call(5, 22, {'region_name': None}, mocks['ServiceQuotasClient'].return_value) ] assert mock_bar.mock_calls == [ call(5, 22, {'region_name': None}, mocks['ServiceQuotasClient'].return_value) ] assert mock_ta_constr.mock_calls == [ call(services, {'region_name': None}, ta_api_region='us-east-1', ta_refresh_mode=None, ta_refresh_timeout=None) ] assert mock_svc1.mock_calls == [] assert mock_svc2.mock_calls == [] assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info def test_init_region_profile_role_partition_ta_region(self): mock_svc1 = Mock(spec_set=_AwsService) mock_svc2 = Mock(spec_set=_AwsService) mock_foo = Mock(spec_set=_AwsService) mock_bar = Mock(spec_set=_AwsService) mock_ta = Mock(spec_set=TrustedAdvisor) mock_foo.return_value = mock_svc1 mock_bar.return_value = mock_svc2 svcs = {'SvcFoo': mock_foo, 'SvcBar': mock_bar} with patch('%s.boto3' % pbm) as mock_boto3: with patch.dict('%s._services' % pbm, values=svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, _get_latest_version=DEFAULT, autospec=True, ) as mocks: mock_boto3.Session.return_value._session = Mock() mock_version = mocks['_get_version_info'] mock_version.return_value = self.mock_ver_info mocks['TrustedAdvisor'].return_value = mock_ta mocks['_get_latest_version'].return_value = None with patch('%s._boto_conn_kwargs' % pb, new_callable=PropertyMock) as m_bck: m_bck.return_value = {'region_name': 'rName'} cls = AwsLimitChecker(region='regionX', profile_name='foo', role_partition='rpName', ta_api_region='taRegion') # dict should be of _AwsService instances services = {'SvcFoo': mock_svc1, 'SvcBar': mock_svc2} assert cls.profile_name == 'foo' assert cls.region == 'regionX' assert cls.services == services assert cls.role_partition == 'rpName' assert mock_svc1.mock_calls == [] assert mock_svc2.mock_calls == [] assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info assert mocks['TrustedAdvisor'].mock_calls == [ call(services, {'region_name': 'rName'}, ta_api_region='taRegion', ta_refresh_mode=None, ta_refresh_timeout=None) ] def test_init_sts(self): mock_svc1 = Mock(spec_set=_AwsService) mock_svc2 = Mock(spec_set=_AwsService) mock_foo = Mock(spec_set=_AwsService) mock_bar = Mock(spec_set=_AwsService) mock_ta = Mock(spec_set=TrustedAdvisor) mock_foo.return_value = mock_svc1 mock_bar.return_value = mock_svc2 svcs = {'SvcFoo': mock_foo, 'SvcBar': mock_bar} with patch('%s.boto3' % pbm) as mock_boto: mock_boto.client.return_value.assume_role.return_value = { 'Credentials': { 'AccessKeyId': 'akid', 'SecretAccessKey': 'sk', 'SessionToken': 'stoken', 'Expiration': '0' }, 'AssumedRoleUser': { 'AssumedRoleId': 'arid', 'Arn': 'arn' } } with patch.dict('%s._services' % pbm, values=svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, _get_latest_version=DEFAULT, autospec=True, ) as mocks: mock_version = mocks['_get_version_info'] mock_version.return_value = self.mock_ver_info mocks['TrustedAdvisor'].return_value = mock_ta mocks['_get_latest_version'].return_value = None cls = AwsLimitChecker(account_id='123456789012', account_role='myrole', region='myregion') # dict should be of _AwsService instances services = {'SvcFoo': mock_svc1, 'SvcBar': mock_svc2} assert cls.services == services assert mock_svc1.mock_calls == [] assert mock_svc2.mock_calls == [] assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info assert mock_boto.mock_calls == [ call.client('sts', region_name='myregion'), call.client().assume_role( RoleArn='arn:aws:iam::123456789012:role/myrole', RoleSessionName='awslimitchecker') ] def test_init_sts_external_id_ta_refresh(self): mock_svc1 = Mock(spec_set=_AwsService) mock_svc2 = Mock(spec_set=_AwsService) mock_foo = Mock(spec_set=_AwsService) mock_bar = Mock(spec_set=_AwsService) mock_ta = Mock(spec_set=TrustedAdvisor) mock_foo.return_value = mock_svc1 mock_bar.return_value = mock_svc2 svcs = {'SvcFoo': mock_foo, 'SvcBar': mock_bar} with patch('%s.boto3' % pbm) as mock_boto: mock_boto.client.return_value.assume_role.return_value = { 'Credentials': { 'AccessKeyId': 'akid', 'SecretAccessKey': 'sk', 'SessionToken': 'stoken', 'Expiration': '0' }, 'AssumedRoleUser': { 'AssumedRoleId': 'arid', 'Arn': 'arn' } } with patch.dict('%s._services' % pbm, values=svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, _get_latest_version=DEFAULT, autospec=True, ) as mocks: mock_version = mocks['_get_version_info'] mock_version.return_value = self.mock_ver_info mocks['TrustedAdvisor'].return_value = mock_ta mocks['_get_latest_version'].return_value = None cls = AwsLimitChecker(account_id='123456789012', account_role='myrole', region='myregion', external_id='myextid', mfa_serial_number=123, mfa_token=456, ta_refresh_mode=123, ta_refresh_timeout=456, role_partition='mypart') # dict should be of _AwsService instances services = {'SvcFoo': mock_svc1, 'SvcBar': mock_svc2} assert cls.services == services assert mock_svc1.mock_calls == [] assert mock_svc2.mock_calls == [] assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info assert mock_boto.mock_calls == [ call.client('sts', region_name='myregion'), call.client().assume_role( ExternalId='myextid', RoleArn='arn:mypart:iam::123456789012:role/myrole', RoleSessionName='awslimitchecker', SerialNumber=123, TokenCode=456) ] def test_boto3_connection_kwargs(self): cls = AwsLimitChecker() with patch('%s._get_sts_token' % pb) as mock_get_sts: with patch('%s.logger' % pbm) as mock_logger: with patch('%s.boto3.Session' % pbm) as mock_sess: res = cls._boto_conn_kwargs assert mock_get_sts.mock_calls == [] assert mock_logger.mock_calls == [ call.debug('Connecting to region %s', None) ] assert mock_sess.mock_calls == [] assert res == {'region_name': None} def test_boto3_connection_kwargs_profile(self): m_creds = Mock() type(m_creds).access_key = 'ak' type(m_creds).secret_key = 'sk' type(m_creds).token = 'tkn' mock_session = Mock() m_sess = Mock() m_sess.get_credentials.return_value = m_creds type(mock_session)._session = m_sess with patch('%s._get_sts_token' % pb) as mock_get_sts: with patch('%s.logger' % pbm) as mock_logger: with patch('%s.boto3.Session' % pbm) as mock_sess: with patch.dict('%s._services' % pbm, {}, clear=True): mock_sess.return_value = mock_session cls = AwsLimitChecker(profile_name='myprof') mock_get_sts.reset_mock() mock_logger.reset_mock() mock_sess.reset_mock() res = cls._boto_conn_kwargs assert mock_get_sts.mock_calls == [] assert mock_logger.mock_calls == [ call.debug('Using credentials profile: %s', 'myprof') ] assert mock_sess.mock_calls == [call(profile_name='myprof')] assert res == { 'region_name': None, 'aws_access_key_id': 'ak', 'aws_secret_access_key': 'sk', 'aws_session_token': 'tkn' } def test_boto3_connection_kwargs_region(self): with patch('%s.boto3' % pbm): cls = AwsLimitChecker(region='myregion') with patch('%s._get_sts_token' % pb) as mock_get_sts: with patch('%s.logger' % pbm) as mock_logger: with patch('%s.boto3.Session' % pbm) as mock_sess: res = cls._boto_conn_kwargs assert mock_get_sts.mock_calls == [] assert mock_logger.mock_calls == [ call.debug('Connecting to region %s', 'myregion') ] assert mock_sess.mock_calls == [] assert res == {'region_name': 'myregion'} def test_boto3_connection_kwargs_sts(self): mock_creds = Mock() type(mock_creds).access_key = 'sts_ak' type(mock_creds).secret_key = 'sts_sk' type(mock_creds).session_token = 'sts_token' with patch('%s._get_sts_token' % pb) as mock_get_sts: with patch('%s.logger' % pbm) as mock_logger: with patch('%s.boto3.Session' % pbm) as mock_sess: with patch.dict('%s._services' % pbm, {}, clear=True): cls = AwsLimitChecker(account_id='123', account_role='myrole', region='myregion') mock_get_sts.return_value = mock_creds mock_get_sts.reset_mock() mock_logger.reset_mock() mock_sess.reset_mock() res = cls._boto_conn_kwargs assert mock_get_sts.mock_calls == [call()] assert mock_logger.mock_calls == [ call.debug( "Connecting for account %s role '%s' with STS " "(region: %s)", '123', 'myrole', 'myregion') ] assert mock_sess.mock_calls == [] assert res == { 'region_name': 'myregion', 'aws_access_key_id': 'sts_ak', 'aws_secret_access_key': 'sts_sk', 'aws_session_token': 'sts_token' } def test_get_version(self): with patch('%s._get_version_info' % pbm, spec_set=_get_version_info) as mock_version: self.cls.vinfo = self.mock_ver_info res = self.cls.get_version() assert res == '1.2.3@mytag' assert mock_version.mock_calls == [] def test_get_project_url(self): with patch('%s._get_version_info' % pbm, spec_set=_get_version_info) as mock_version: self.cls.vinfo = self.mock_ver_info res = self.cls.get_project_url() assert res == 'http://myurl' assert mock_version.mock_calls == [] def test_remove_services_none(self): self.cls.remove_services() assert self.cls.services == { 'SvcFoo': self.mock_svc1, 'SvcBar': self.mock_svc2 } def test_remove_services_one(self): self.cls.remove_services(['SvcFoo']) assert self.cls.services == {'SvcBar': self.mock_svc2} def test_remove_services_all(self): self.cls.remove_services(['SvcFoo', 'SvcBar']) assert self.cls.services == {} def test_get_service_names(self): res = self.cls.get_service_names() assert res == ['SvcBar', 'SvcFoo'] def test_get_limits(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] res = self.cls.get_limits() assert res == limits assert self.mock_ta.mock_calls == [call.update_limits()] assert self.mock_svc1.mock_calls == [ call._update_service_quotas(), call.get_limits() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call._update_service_quotas(), call.get_limits() ] def test_get_limits_no_ta(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] res = self.cls.get_limits(use_ta=False) assert res == limits assert self.mock_ta.mock_calls == [] assert self.mock_svc1.mock_calls == [ call._update_service_quotas(), call.get_limits() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call._update_service_quotas(), call.get_limits() ] def test_get_limits_service(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] res = self.cls.get_limits(service=['SvcFoo']) assert res == {'SvcFoo': limits['SvcFoo']} assert self.mock_ta.mock_calls == [call.update_limits()] assert self.mock_svc1.mock_calls == [ call._update_service_quotas(), call.get_limits() ] assert self.mock_svc2.mock_calls == [] def test_get_limits_service_with_api(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] res = self.cls.get_limits(service=['SvcBar']) assert res == {'SvcBar': limits['SvcBar']} assert self.mock_ta.mock_calls == [call.update_limits()] assert self.mock_svc1.mock_calls == [] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call._update_service_quotas(), call.get_limits() ] def test_find_usage(self): self.cls.find_usage() assert self.mock_svc1.mock_calls == [ call._update_service_quotas(), call.find_usage() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call._update_service_quotas(), call.find_usage() ] assert self.mock_ta.mock_calls == [call.update_limits()] def test_find_usage_no_ta(self): self.cls.find_usage(use_ta=False) assert self.mock_svc1.mock_calls == [ call._update_service_quotas(), call.find_usage() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call._update_service_quotas(), call.find_usage() ] assert self.mock_ta.mock_calls == [] def test_find_usage_service(self): self.cls.find_usage(service=['SvcFoo']) assert self.mock_svc1.mock_calls == [ call._update_service_quotas(), call.find_usage() ] assert self.mock_svc2.mock_calls == [] assert self.mock_ta.mock_calls == [call.update_limits()] def test_find_usage_service_with_api(self): self.cls.find_usage(service=['SvcBar']) assert self.mock_svc1.mock_calls == [] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call._update_service_quotas(), call.find_usage() ] assert self.mock_ta.mock_calls == [call.update_limits()] def test_set_threshold_overrides(self): limits = sample_limits() limits['SvcFoo']['zz3'] = AwsLimit( 'zz3', self.mock_svc1, 1, 2, 3, ) self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] overrides = { 'SvcBar': { 'barlimit1': { 'warning': { 'percent': 10, 'count': 12 }, 'critical': { 'percent': 14, 'count': 16 } }, 'bar limit2': { 'critical': { 'count': 15, } }, 'zz3': { 'warning': { 'count': 41 }, 'critical': { 'percent': 52 } } }, 'SvcFoo': { 'foo limit3': { 'warning': { 'percent': 91 }, } }, } self.cls.set_threshold_overrides(overrides) assert self.mock_svc1.mock_calls == [ call.set_threshold_override( 'foo limit3', warn_percent=91, ) ] assert self.mock_svc2.mock_calls == [ call.set_threshold_override('bar limit2', crit_count=15), call.set_threshold_override('barlimit1', warn_percent=10, warn_count=12, crit_percent=14, crit_count=16), call.set_threshold_override('zz3', warn_count=41, crit_percent=52), ] def test_set_limit_overrides(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] overrides = { 'SvcBar': { 'barlimit1': 100, }, 'SvcFoo': { 'foo limit3': 99, }, } self.cls.set_limit_overrides(overrides) assert self.mock_svc1.mock_calls == [ call.set_limit_override('foo limit3', 99, override_ta=True) ] assert self.mock_svc2.mock_calls == [ call.set_limit_override('barlimit1', 100, override_ta=True) ] def test_set_limit_overrides_ta(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] overrides = { 'SvcBar': { 'bar limit2': 100, }, 'SvcFoo': { 'foo limit3': 3, }, } self.cls.set_limit_overrides(overrides, override_ta=False) assert self.mock_svc1.mock_calls == [ call.set_limit_override('foo limit3', 3, override_ta=False) ] assert self.mock_svc2.mock_calls == [ call.set_limit_override('bar limit2', 100, override_ta=False) ] def test_set_limit_override(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.cls.set_limit_override('SvcFoo', 'foo limit3', 99) assert self.mock_svc1.mock_calls == [ call.set_limit_override('foo limit3', 99, override_ta=True) ] def test_set_limit_override_ta(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.cls.set_limit_override('SvcFoo', 'foo limit3', 99, override_ta=False) assert self.mock_svc1.mock_calls == [ call.set_limit_override('foo limit3', 99, override_ta=False) ] def test_set_threshold_override(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.cls.set_threshold_override('SvcFoo', 'foo limit3', warn_percent=10, warn_count=12, crit_percent=14, crit_count=16) assert self.mock_svc1.mock_calls == [ call.set_threshold_override('foo limit3', warn_percent=10, warn_count=12, crit_percent=14, crit_count=16) ] def test_get_required_iam_policy(self): expected = { 'Version': '2012-10-17', 'Statement': [{ 'Effect': 'Allow', 'Resource': '*', 'Action': [ 'cloudwatch:GetMetricData', 'ec2:bar', 'ec2:foo', 'foo:perm1', 'foo:perm2', 'servicequotas:ListServiceQuotas', 'sts:GetCallerIdentity', 'support:DescribeTrustedAdvisorCheckRefreshStatuses', 'support:DescribeTrustedAdvisorCheckResult', 'support:DescribeTrustedAdvisorCheckSummaries', 'support:DescribeTrustedAdvisorChecks', 'support:RefreshTrustedAdvisorCheck', 'trustedadvisor:Describe*', 'trustedadvisor:RefreshCheck' ], }], } self.mock_svc1.required_iam_permissions.return_value = [ 'ec2:foo', 'ec2:bar', 'foo:perm1' ] self.mock_svc2.required_iam_permissions.return_value = [ 'foo:perm1', 'foo:perm2', ] res = self.cls.get_required_iam_policy() assert res == expected assert self.mock_svc1.mock_calls == [call.required_iam_permissions()] assert self.mock_svc2.mock_calls == [call.required_iam_permissions()] def test_check_thresholds(self): self.mock_svc1.check_thresholds.return_value = { 'foo': 'bar', 'baz': 'blam', } self.mock_svc2.check_thresholds.return_value = {} res = self.cls.check_thresholds() assert res == { 'SvcFoo': { 'foo': 'bar', 'baz': 'blam', } } assert self.mock_ta.mock_calls == [ call.update_limits(), ] assert self.mock_svc1.mock_calls == [ call._update_service_quotas(), call.check_thresholds() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call._update_service_quotas(), call.check_thresholds() ] def test_check_thresholds_service(self): self.mock_svc1.check_thresholds.return_value = {'foo': 'bar'} self.mock_svc2.check_thresholds.return_value = {'baz': 'blam'} res = self.cls.check_thresholds(service=['SvcFoo']) assert res == { 'SvcFoo': { 'foo': 'bar', } } assert self.mock_ta.mock_calls == [call.update_limits()] assert self.mock_svc1.mock_calls == [ call._update_service_quotas(), call.check_thresholds() ] assert self.mock_svc2.mock_calls == [] def test_check_thresholds_service_api(self): self.mock_svc1.check_thresholds.return_value = {'foo': 'bar'} self.mock_svc2.check_thresholds.return_value = {'baz': 'blam'} res = self.cls.check_thresholds(service=['SvcBar']) assert res == { 'SvcBar': { 'baz': 'blam', } } assert self.mock_ta.mock_calls == [call.update_limits()] assert self.mock_svc1.mock_calls == [] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call._update_service_quotas(), call.check_thresholds() ] def test_check_thresholds_no_ta(self): self.mock_svc1.check_thresholds.return_value = { 'foo': 'bar', 'baz': 'blam', } self.mock_svc2.check_thresholds.return_value = {} self.cls.use_ta = False res = self.cls.check_thresholds(use_ta=False) assert res == { 'SvcFoo': { 'foo': 'bar', 'baz': 'blam', } } assert self.mock_ta.mock_calls == [] assert self.mock_svc1.mock_calls == [ call._update_service_quotas(), call.check_thresholds() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call._update_service_quotas(), call.check_thresholds() ] def test_region_name(self): mock_client = Mock(_client_config=Mock(region_name='rname')) with patch('%s._boto_conn_kwargs' % pb, new_callable=PropertyMock) as mock_bck: mock_bck.return_value = {'foo': 'bar'} with patch('%s.boto3.client' % pbm) as m_client: m_client.return_value = mock_client res = self.cls.region_name assert res == 'rname'
def check_limits(self, acct_id, region_name, role_name=None, limit_overrides={}, threshold_overrides={}): """ Run the actual usage and limit check, with overrides, against a specific account in a specific region, optionally assuming a role in the account and optionally setting limit and/or threshold overrides. Return a 2-tuple of lists, warning strings and critical strings. see: http://awslimitchecker.readthedocs.org/en/latest/python_usage.html :returns: 2-tuple of lists of strings, warnings and criticals """ # instantiate the class if role_name is not None: checker = AwsLimitChecker(account_id=acct_id, region=region_name, account_role=role_name) else: checker = AwsLimitChecker(region=region_name) # set your overrides if len(threshold_overrides) > 0: checker.set_threshold_overrides(threshold_overrides) if len(limit_overrides) > 0: checker.set_limit_overrides(limit_overrides) # check usage against thresholds checker.check_thresholds() # save state for exit code and summary warnings = [] criticals = [] # iterate the results for service, svc_limits in sorted(checker.get_limits().items()): for limit_name, limit in sorted(svc_limits.items()): # check warnings and criticals for each Limit for warn in limit.get_warnings(): warnings.append( colored( "{service} '{limit_name}' usage " "({u}) exceeds warning threshold " "(limit={l})".format( service=service, limit_name=limit_name, u=str(warn), l=limit.get_limit(), ), 'yellow')) for crit in limit.get_criticals(): criticals.append( colored( "{service} '{limit_name}' usage " "({u}) exceeds critical threshold" " (limit={l})".format( service=service, limit_name=limit_name, u=str(crit), l=limit.get_limit(), ), 'red')) return warnings, criticals
class TestAwsLimitChecker(object): def setup(self): self.mock_ver_info = Mock( release='1.2.3', url='http://myurl', commit='abcd', tag='mytag', version_str='1.2.3@mytag' ) self.mock_svc1 = Mock(spec_set=_AwsService) self.mock_svc2 = Mock(spec_set=ApiServiceSpec) self.mock_foo = Mock(spec_set=_AwsService) self.mock_bar = Mock(spec_set=_AwsService) self.mock_ta = Mock(spec_set=TrustedAdvisor) self.mock_foo.return_value = self.mock_svc1 self.mock_bar.return_value = self.mock_svc2 self.svcs = {'SvcFoo': self.mock_foo, 'SvcBar': self.mock_bar} with patch.dict('%s._services' % pbm, values=self.svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, autospec=True, ) as mocks: self.mock_logger = mocks['logger'] self.mock_version = mocks['_get_version_info'] self.mock_ta_constr = mocks['TrustedAdvisor'] mocks['TrustedAdvisor'].return_value = self.mock_ta self.mock_version.return_value = self.mock_ver_info self.cls = AwsLimitChecker() def test_init(self): # dict should be of _AwsService instances services = { 'SvcFoo': self.mock_svc1, 'SvcBar': self.mock_svc2 } assert self.cls.services == services # _AwsService instances should exist, but have no other calls assert self.mock_foo.mock_calls == [ call(80, 99, {'region_name': None}) ] assert self.mock_bar.mock_calls == [ call(80, 99, {'region_name': None}) ] assert self.mock_ta_constr.mock_calls == [ call(services, {'region_name': None}, ta_refresh_mode=None, ta_refresh_timeout=None) ] assert self.mock_svc1.mock_calls == [] assert self.mock_svc2.mock_calls == [] assert self.cls.ta == self.mock_ta assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info assert self.mock_logger.mock_calls == [ call.debug('Connecting to region %s', None) ] def test_init_AGPL_message(self, capsys): # get rid of the class self.cls = None # clear out/err out, err = capsys.readouterr() # run setup again self.setup() # check out/err out, err = capsys.readouterr() assert out == '' assert (err) == ( "awslimitchecker 1.2.3@mytag is AGPL-licensed free software; " "all users have a right to the full source code of " "this version. See <http://myurl>\n") def test_init_thresholds(self): mock_svc1 = Mock(spec_set=_AwsService) mock_svc2 = Mock(spec_set=_AwsService) mock_foo = Mock(spec_set=_AwsService) mock_bar = Mock(spec_set=_AwsService) mock_ta = Mock(spec_set=TrustedAdvisor) mock_foo.return_value = mock_svc1 mock_bar.return_value = mock_svc2 svcs = {'SvcFoo': mock_foo, 'SvcBar': mock_bar} with patch.dict('%s._services' % pbm, values=svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, autospec=True, ) as mocks: mock_version = mocks['_get_version_info'] mock_version.return_value = self.mock_ver_info mock_ta_constr = mocks['TrustedAdvisor'] mocks['TrustedAdvisor'].return_value = mock_ta cls = AwsLimitChecker( warning_threshold=5, critical_threshold=22, ) # dict should be of _AwsService instances services = { 'SvcFoo': mock_svc1, 'SvcBar': mock_svc2 } assert cls.services == services # _AwsService instances should exist, but have no other calls assert mock_foo.mock_calls == [ call(5, 22, {'region_name': None}) ] assert mock_bar.mock_calls == [ call(5, 22, {'region_name': None}) ] assert mock_ta_constr.mock_calls == [ call(services, {'region_name': None}, ta_refresh_mode=None, ta_refresh_timeout=None) ] assert mock_svc1.mock_calls == [] assert mock_svc2.mock_calls == [] assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info def test_init_region_profile(self): mock_svc1 = Mock(spec_set=_AwsService) mock_svc2 = Mock(spec_set=_AwsService) mock_foo = Mock(spec_set=_AwsService) mock_bar = Mock(spec_set=_AwsService) mock_ta = Mock(spec_set=TrustedAdvisor) mock_foo.return_value = mock_svc1 mock_bar.return_value = mock_svc2 svcs = {'SvcFoo': mock_foo, 'SvcBar': mock_bar} with patch('%s.boto3' % pbm) as mock_boto3: with patch.dict('%s._services' % pbm, values=svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, autospec=True, ) as mocks: mock_boto3.Session.return_value._session = Mock() mock_version = mocks['_get_version_info'] mock_version.return_value = self.mock_ver_info mocks['TrustedAdvisor'].return_value = mock_ta cls = AwsLimitChecker(region='regionX', profile_name='foo') # dict should be of _AwsService instances services = { 'SvcFoo': mock_svc1, 'SvcBar': mock_svc2 } assert cls.profile_name == 'foo' assert cls.region == 'regionX' assert cls.services == services assert mock_svc1.mock_calls == [] assert mock_svc2.mock_calls == [] assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info def test_init_sts(self): mock_svc1 = Mock(spec_set=_AwsService) mock_svc2 = Mock(spec_set=_AwsService) mock_foo = Mock(spec_set=_AwsService) mock_bar = Mock(spec_set=_AwsService) mock_ta = Mock(spec_set=TrustedAdvisor) mock_foo.return_value = mock_svc1 mock_bar.return_value = mock_svc2 svcs = {'SvcFoo': mock_foo, 'SvcBar': mock_bar} with patch('%s.boto3' % pbm): with patch.dict('%s._services' % pbm, values=svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, autospec=True, ) as mocks: mock_version = mocks['_get_version_info'] mock_version.return_value = self.mock_ver_info mocks['TrustedAdvisor'].return_value = mock_ta cls = AwsLimitChecker( account_id='123456789012', account_role='myrole', region='myregion' ) # dict should be of _AwsService instances services = { 'SvcFoo': mock_svc1, 'SvcBar': mock_svc2 } assert cls.services == services assert mock_svc1.mock_calls == [] assert mock_svc2.mock_calls == [] assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info def test_init_sts_external_id_ta_refresh(self): mock_svc1 = Mock(spec_set=_AwsService) mock_svc2 = Mock(spec_set=_AwsService) mock_foo = Mock(spec_set=_AwsService) mock_bar = Mock(spec_set=_AwsService) mock_ta = Mock(spec_set=TrustedAdvisor) mock_foo.return_value = mock_svc1 mock_bar.return_value = mock_svc2 svcs = {'SvcFoo': mock_foo, 'SvcBar': mock_bar} with patch('%s.boto3' % pbm): with patch.dict('%s._services' % pbm, values=svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, autospec=True, ) as mocks: mock_version = mocks['_get_version_info'] mock_version.return_value = self.mock_ver_info mocks['TrustedAdvisor'].return_value = mock_ta cls = AwsLimitChecker( account_id='123456789012', account_role='myrole', region='myregion', external_id='myextid', mfa_serial_number=123, mfa_token=456, ta_refresh_mode=123, ta_refresh_timeout=456 ) # dict should be of _AwsService instances services = { 'SvcFoo': mock_svc1, 'SvcBar': mock_svc2 } assert cls.services == services assert mock_svc1.mock_calls == [] assert mock_svc2.mock_calls == [] assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info def test_boto3_connection_kwargs(self): cls = AwsLimitChecker() with patch('%s._get_sts_token' % pb) as mock_get_sts: with patch('%s.logger' % pbm) as mock_logger: with patch('%s.boto3.Session' % pbm) as mock_sess: res = cls._boto_conn_kwargs assert mock_get_sts.mock_calls == [] assert mock_logger.mock_calls == [ call.debug('Connecting to region %s', None) ] assert mock_sess.mock_calls == [] assert res == { 'region_name': None } def test_boto3_connection_kwargs_profile(self): with patch('%s.boto3' % pbm): cls = AwsLimitChecker(profile_name='myprof') m_creds = Mock() type(m_creds).access_key = 'ak' type(m_creds).secret_key = 'sk' type(m_creds).token = 'tkn' mock_session = Mock() m_sess = Mock() m_sess.get_credentials.return_value = m_creds type(mock_session)._session = m_sess with patch('%s._get_sts_token' % pb) as mock_get_sts: with patch('%s.logger' % pbm) as mock_logger: with patch('%s.boto3.Session' % pbm) as mock_sess: mock_sess.return_value = mock_session res = cls._boto_conn_kwargs assert mock_get_sts.mock_calls == [] assert mock_logger.mock_calls == [ call.debug('Using credentials profile: %s', 'myprof') ] assert mock_sess.mock_calls == [call(profile_name='myprof')] assert res == { 'region_name': None, 'aws_access_key_id': 'ak', 'aws_secret_access_key': 'sk', 'aws_session_token': 'tkn' } def test_boto3_connection_kwargs_region(self): with patch('%s.boto3' % pbm): cls = AwsLimitChecker(region='myregion') with patch('%s._get_sts_token' % pb) as mock_get_sts: with patch('%s.logger' % pbm) as mock_logger: with patch('%s.boto3.Session' % pbm) as mock_sess: res = cls._boto_conn_kwargs assert mock_get_sts.mock_calls == [] assert mock_logger.mock_calls == [ call.debug('Connecting to region %s', 'myregion') ] assert mock_sess.mock_calls == [] assert res == { 'region_name': 'myregion' } def test_boto3_connection_kwargs_sts(self): with patch('%s.boto3' % pbm): cls = AwsLimitChecker(account_id='123', account_role='myrole', region='myregion') mock_creds = Mock() type(mock_creds).access_key = 'sts_ak' type(mock_creds).secret_key = 'sts_sk' type(mock_creds).session_token = 'sts_token' with patch('%s._get_sts_token' % pb) as mock_get_sts: with patch('%s.logger' % pbm) as mock_logger: with patch('%s.boto3.Session' % pbm) as mock_sess: mock_get_sts.return_value = mock_creds res = cls._boto_conn_kwargs assert mock_get_sts.mock_calls == [call()] assert mock_logger.mock_calls == [ call.debug("Connecting for account %s role '%s' with STS " "(region: %s)", '123', 'myrole', 'myregion') ] assert mock_sess.mock_calls == [] assert res == { 'region_name': 'myregion', 'aws_access_key_id': 'sts_ak', 'aws_secret_access_key': 'sts_sk', 'aws_session_token': 'sts_token' } def test_get_version(self): with patch('%s._get_version_info' % pbm, spec_set=_get_version_info) as mock_version: self.cls.vinfo = self.mock_ver_info res = self.cls.get_version() assert res == '1.2.3@mytag' assert mock_version.mock_calls == [] def test_get_project_url(self): with patch('%s._get_version_info' % pbm, spec_set=_get_version_info) as mock_version: self.cls.vinfo = self.mock_ver_info res = self.cls.get_project_url() assert res == 'http://myurl' assert mock_version.mock_calls == [] def test_get_service_names(self): res = self.cls.get_service_names() assert res == ['SvcBar', 'SvcFoo'] def test_get_limits(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] res = self.cls.get_limits() assert res == limits assert self.mock_ta.mock_calls == [ call.update_limits() ] assert self.mock_svc1.mock_calls == [ call.get_limits() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.get_limits() ] def test_get_limits_no_ta(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] res = self.cls.get_limits(use_ta=False) assert res == limits assert self.mock_ta.mock_calls == [] assert self.mock_svc1.mock_calls == [ call.get_limits() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.get_limits() ] def test_get_limits_service(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] res = self.cls.get_limits(service=['SvcFoo']) assert res == {'SvcFoo': limits['SvcFoo']} assert self.mock_ta.mock_calls == [ call.update_limits() ] assert self.mock_svc1.mock_calls == [ call.get_limits() ] assert self.mock_svc2.mock_calls == [] def test_get_limits_service_with_api(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] res = self.cls.get_limits(service=['SvcBar']) assert res == {'SvcBar': limits['SvcBar']} assert self.mock_ta.mock_calls == [ call.update_limits() ] assert self.mock_svc1.mock_calls == [] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.get_limits() ] def test_find_usage(self): self.cls.find_usage() assert self.mock_svc1.mock_calls == [ call.find_usage() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.find_usage() ] assert self.mock_ta.mock_calls == [ call.update_limits() ] def test_find_usage_no_ta(self): self.cls.find_usage(use_ta=False) assert self.mock_svc1.mock_calls == [ call.find_usage() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.find_usage() ] assert self.mock_ta.mock_calls == [] def test_find_usage_service(self): self.cls.find_usage(service=['SvcFoo']) assert self.mock_svc1.mock_calls == [ call.find_usage() ] assert self.mock_svc2.mock_calls == [] assert self.mock_ta.mock_calls == [ call.update_limits() ] def test_find_usage_service_with_api(self): self.cls.find_usage(service=['SvcBar']) assert self.mock_svc1.mock_calls == [] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.find_usage() ] assert self.mock_ta.mock_calls == [ call.update_limits() ] def test_set_threshold_overrides(self): limits = sample_limits() limits['SvcFoo']['zz3'] = AwsLimit( 'zz3', self.mock_svc1, 1, 2, 3, ) self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] overrides = { 'SvcBar': { 'barlimit1': { 'warning': { 'percent': 10, 'count': 12 }, 'critical': { 'percent': 14, 'count': 16 } }, 'bar limit2': { 'critical': { 'count': 15, } }, 'zz3': { 'warning': { 'count': 41 }, 'critical': { 'percent': 52 } } }, 'SvcFoo': { 'foo limit3': { 'warning': { 'percent': 91 }, } }, } self.cls.set_threshold_overrides(overrides) assert self.mock_svc1.mock_calls == [ call.set_threshold_override( 'foo limit3', warn_percent=91, ) ] assert self.mock_svc2.mock_calls == [ call.set_threshold_override( 'bar limit2', crit_count=15 ), call.set_threshold_override( 'barlimit1', warn_percent=10, warn_count=12, crit_percent=14, crit_count=16 ), call.set_threshold_override( 'zz3', warn_count=41, crit_percent=52 ), ] def test_set_limit_overrides(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] overrides = { 'SvcBar': { 'barlimit1': 100, }, 'SvcFoo': { 'foo limit3': 99, }, } self.cls.set_limit_overrides(overrides) assert self.mock_svc1.mock_calls == [ call.set_limit_override('foo limit3', 99, override_ta=True) ] assert self.mock_svc2.mock_calls == [ call.set_limit_override('barlimit1', 100, override_ta=True) ] def test_set_limit_overrides_ta(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] overrides = { 'SvcBar': { 'bar limit2': 100, }, 'SvcFoo': { 'foo limit3': 3, }, } self.cls.set_limit_overrides(overrides, override_ta=False) assert self.mock_svc1.mock_calls == [ call.set_limit_override('foo limit3', 3, override_ta=False) ] assert self.mock_svc2.mock_calls == [ call.set_limit_override('bar limit2', 100, override_ta=False) ] def test_set_limit_override(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.cls.set_limit_override('SvcFoo', 'foo limit3', 99) assert self.mock_svc1.mock_calls == [ call.set_limit_override('foo limit3', 99, override_ta=True) ] def test_set_limit_override_ta(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.cls.set_limit_override( 'SvcFoo', 'foo limit3', 99, override_ta=False ) assert self.mock_svc1.mock_calls == [ call.set_limit_override( 'foo limit3', 99, override_ta=False ) ] def test_set_threshold_override(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.cls.set_threshold_override( 'SvcFoo', 'foo limit3', warn_percent=10, warn_count=12, crit_percent=14, crit_count=16 ) assert self.mock_svc1.mock_calls == [ call.set_threshold_override( 'foo limit3', warn_percent=10, warn_count=12, crit_percent=14, crit_count=16 ) ] def test_get_required_iam_policy(self): expected = { 'Version': '2012-10-17', 'Statement': [{ 'Effect': 'Allow', 'Resource': '*', 'Action': [ 'ec2:bar', 'ec2:foo', 'foo:perm1', 'foo:perm2', 'support:*', 'trustedadvisor:Describe*', 'trustedadvisor:RefreshCheck' ], }], } self.mock_svc1.required_iam_permissions.return_value = [ 'ec2:foo', 'ec2:bar', 'foo:perm1' ] self.mock_svc2.required_iam_permissions.return_value = [ 'foo:perm1', 'foo:perm2', ] res = self.cls.get_required_iam_policy() assert res == expected assert self.mock_svc1.mock_calls == [call.required_iam_permissions()] assert self.mock_svc2.mock_calls == [call.required_iam_permissions()] def test_check_thresholds(self): self.mock_svc1.check_thresholds.return_value = { 'foo': 'bar', 'baz': 'blam', } self.mock_svc2.check_thresholds.return_value = {} res = self.cls.check_thresholds() assert res == { 'SvcFoo': { 'foo': 'bar', 'baz': 'blam', } } assert self.mock_ta.mock_calls == [ call.update_limits(), ] assert self.mock_svc1.mock_calls == [ call.check_thresholds() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.check_thresholds() ] def test_check_thresholds_service(self): self.mock_svc1.check_thresholds.return_value = {'foo': 'bar'} self.mock_svc2.check_thresholds.return_value = {'baz': 'blam'} res = self.cls.check_thresholds(service=['SvcFoo']) assert res == { 'SvcFoo': { 'foo': 'bar', } } assert self.mock_ta.mock_calls == [ call.update_limits() ] assert self.mock_svc1.mock_calls == [ call.check_thresholds() ] assert self.mock_svc2.mock_calls == [] def test_check_thresholds_service_api(self): self.mock_svc1.check_thresholds.return_value = {'foo': 'bar'} self.mock_svc2.check_thresholds.return_value = {'baz': 'blam'} res = self.cls.check_thresholds(service=['SvcBar']) assert res == { 'SvcBar': { 'baz': 'blam', } } assert self.mock_ta.mock_calls == [ call.update_limits() ] assert self.mock_svc1.mock_calls == [] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.check_thresholds() ] def test_check_thresholds_no_ta(self): self.mock_svc1.check_thresholds.return_value = { 'foo': 'bar', 'baz': 'blam', } self.mock_svc2.check_thresholds.return_value = {} self.cls.use_ta = False res = self.cls.check_thresholds(use_ta=False) assert res == { 'SvcFoo': { 'foo': 'bar', 'baz': 'blam', } } assert self.mock_ta.mock_calls == [] assert self.mock_svc1.mock_calls == [ call.check_thresholds() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.check_thresholds() ]
class TestAwsLimitChecker(object): def setup(self): self.mock_ver_info = Mock( release='1.2.3', url='http://myurl', commit='abcd', tag='mytag', version_str='1.2.3@mytag' ) self.mock_svc1 = Mock(spec_set=_AwsService) self.mock_svc2 = Mock(spec_set=ApiServiceSpec) self.mock_foo = Mock(spec_set=_AwsService) self.mock_bar = Mock(spec_set=_AwsService) self.mock_ta = Mock(spec_set=TrustedAdvisor) self.mock_foo.return_value = self.mock_svc1 self.mock_bar.return_value = self.mock_svc2 self.svcs = {'SvcFoo': self.mock_foo, 'SvcBar': self.mock_bar} with patch.dict('awslimitchecker.checker._services', values=self.svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, autospec=True, ) as mocks: self.mock_logger = mocks['logger'] self.mock_version = mocks['_get_version_info'] self.mock_ta_constr = mocks['TrustedAdvisor'] mocks['TrustedAdvisor'].return_value = self.mock_ta self.mock_version.return_value = self.mock_ver_info self.cls = AwsLimitChecker() def test_init(self): # dict should be of _AwsService instances assert self.cls.services == { 'SvcFoo': self.mock_svc1, 'SvcBar': self.mock_svc2 } # _AwsService instances should exist, but have no other calls assert self.mock_foo.mock_calls == [ call(80, 99, None, None, None, None) ] assert self.mock_bar.mock_calls == [ call(80, 99, None, None, None, None) ] assert self.mock_ta_constr.mock_calls == [ call(account_id=None, account_role=None, region=None, external_id=None) ] assert self.mock_svc1.mock_calls == [] assert self.mock_svc2.mock_calls == [] assert self.cls.ta == self.mock_ta assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info assert self.mock_logger.mock_calls == [] def test_init_AGPL_message(self, capsys): # get rid of the class self.cls = None # clear out/err out, err = capsys.readouterr() # run setup again self.setup() # check out/err out, err = capsys.readouterr() assert out == '' assert (err) == ( "awslimitchecker 1.2.3@mytag is AGPL-licensed free software; " "all users have a right to the full source code of " "this version. See <http://myurl>\n") def test_init_thresholds(self): mock_svc1 = Mock(spec_set=_AwsService) mock_svc2 = Mock(spec_set=_AwsService) mock_foo = Mock(spec_set=_AwsService) mock_bar = Mock(spec_set=_AwsService) mock_ta = Mock(spec_set=TrustedAdvisor) mock_foo.return_value = mock_svc1 mock_bar.return_value = mock_svc2 svcs = {'SvcFoo': mock_foo, 'SvcBar': mock_bar} with patch.dict('awslimitchecker.checker._services', values=svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, autospec=True, ) as mocks: mock_version = mocks['_get_version_info'] mock_version.return_value = self.mock_ver_info mock_ta_constr = mocks['TrustedAdvisor'] mocks['TrustedAdvisor'].return_value = mock_ta cls = AwsLimitChecker( warning_threshold=5, critical_threshold=22, ) # dict should be of _AwsService instances assert cls.services == { 'SvcFoo': mock_svc1, 'SvcBar': mock_svc2 } # _AwsService instances should exist, but have no other calls assert mock_foo.mock_calls == [call(5, 22, None, None, None, None)] assert mock_bar.mock_calls == [call(5, 22, None, None, None, None)] assert mock_ta_constr.mock_calls == [ call(account_id=None, account_role=None, region=None, external_id=None) ] assert mock_svc1.mock_calls == [] assert mock_svc2.mock_calls == [] assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info def test_init_region(self): mock_svc1 = Mock(spec_set=_AwsService) mock_svc2 = Mock(spec_set=_AwsService) mock_foo = Mock(spec_set=_AwsService) mock_bar = Mock(spec_set=_AwsService) mock_ta = Mock(spec_set=TrustedAdvisor) mock_foo.return_value = mock_svc1 mock_bar.return_value = mock_svc2 svcs = {'SvcFoo': mock_foo, 'SvcBar': mock_bar} with patch.dict('awslimitchecker.checker._services', values=svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, autospec=True, ) as mocks: mock_version = mocks['_get_version_info'] mock_version.return_value = self.mock_ver_info mock_ta_constr = mocks['TrustedAdvisor'] mocks['TrustedAdvisor'].return_value = mock_ta cls = AwsLimitChecker(region='myregion') # dict should be of _AwsService instances assert cls.services == { 'SvcFoo': mock_svc1, 'SvcBar': mock_svc2 } # _AwsService instances should exist, but have no other calls assert mock_foo.mock_calls == [ call(80, 99, None, None, 'myregion', None) ] assert mock_bar.mock_calls == [ call(80, 99, None, None, 'myregion', None) ] assert mock_ta_constr.mock_calls == [ call(account_id=None, account_role=None, region='myregion', external_id=None) ] assert mock_svc1.mock_calls == [] assert mock_svc2.mock_calls == [] assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info def test_init_sts(self): mock_svc1 = Mock(spec_set=_AwsService) mock_svc2 = Mock(spec_set=_AwsService) mock_foo = Mock(spec_set=_AwsService) mock_bar = Mock(spec_set=_AwsService) mock_ta = Mock(spec_set=TrustedAdvisor) mock_foo.return_value = mock_svc1 mock_bar.return_value = mock_svc2 svcs = {'SvcFoo': mock_foo, 'SvcBar': mock_bar} with patch.dict('awslimitchecker.checker._services', values=svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, autospec=True, ) as mocks: mock_version = mocks['_get_version_info'] mock_version.return_value = self.mock_ver_info mock_ta_constr = mocks['TrustedAdvisor'] mocks['TrustedAdvisor'].return_value = mock_ta cls = AwsLimitChecker( account_id='123456789012', account_role='myrole', region='myregion' ) # dict should be of _AwsService instances assert cls.services == { 'SvcFoo': mock_svc1, 'SvcBar': mock_svc2 } # _AwsService instances should exist, but have no other calls assert mock_foo.mock_calls == [ call(80, 99, '123456789012', 'myrole', 'myregion', None) ] assert mock_bar.mock_calls == [ call(80, 99, '123456789012', 'myrole', 'myregion', None) ] assert mock_ta_constr.mock_calls == [ call( account_id='123456789012', account_role='myrole', region='myregion', external_id=None ) ] assert mock_svc1.mock_calls == [] assert mock_svc2.mock_calls == [] assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info def test_init_sts_external_id(self): mock_svc1 = Mock(spec_set=_AwsService) mock_svc2 = Mock(spec_set=_AwsService) mock_foo = Mock(spec_set=_AwsService) mock_bar = Mock(spec_set=_AwsService) mock_ta = Mock(spec_set=TrustedAdvisor) mock_foo.return_value = mock_svc1 mock_bar.return_value = mock_svc2 svcs = {'SvcFoo': mock_foo, 'SvcBar': mock_bar} with patch.dict('awslimitchecker.checker._services', values=svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, autospec=True, ) as mocks: mock_version = mocks['_get_version_info'] mock_version.return_value = self.mock_ver_info mock_ta_constr = mocks['TrustedAdvisor'] mocks['TrustedAdvisor'].return_value = mock_ta cls = AwsLimitChecker( account_id='123456789012', account_role='myrole', region='myregion', external_id='myextid' ) # dict should be of _AwsService instances assert cls.services == { 'SvcFoo': mock_svc1, 'SvcBar': mock_svc2 } # _AwsService instances should exist, but have no other calls assert mock_foo.mock_calls == [ call(80, 99, '123456789012', 'myrole', 'myregion', 'myextid') ] assert mock_bar.mock_calls == [ call(80, 99, '123456789012', 'myrole', 'myregion', 'myextid') ] assert mock_ta_constr.mock_calls == [ call( account_id='123456789012', account_role='myrole', region='myregion', external_id='myextid' ) ] assert mock_svc1.mock_calls == [] assert mock_svc2.mock_calls == [] assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info def test_get_version(self): with patch('awslimitchecker.checker._get_version_info', spec_set=_get_version_info) as mock_version: self.cls.vinfo = self.mock_ver_info res = self.cls.get_version() assert res == '1.2.3@mytag' assert mock_version.mock_calls == [] def test_get_project_url(self): with patch('awslimitchecker.checker._get_version_info', spec_set=_get_version_info) as mock_version: self.cls.vinfo = self.mock_ver_info res = self.cls.get_project_url() assert res == 'http://myurl' assert mock_version.mock_calls == [] def test_get_service_names(self): res = self.cls.get_service_names() assert res == ['SvcBar', 'SvcFoo'] def test_get_limits(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] res = self.cls.get_limits() assert res == limits assert self.mock_ta.mock_calls == [ call.update_limits({ 'SvcFoo': self.mock_svc1, 'SvcBar': self.mock_svc2, }) ] assert self.mock_svc1.mock_calls == [ call.get_limits() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.get_limits() ] def test_get_limits_no_ta(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] res = self.cls.get_limits(use_ta=False) assert res == limits assert self.mock_ta.mock_calls == [] assert self.mock_svc1.mock_calls == [ call.get_limits() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.get_limits() ] def test_get_limits_service(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] res = self.cls.get_limits(service='SvcFoo') assert res == {'SvcFoo': limits['SvcFoo']} assert self.mock_ta.mock_calls == [ call.update_limits({ 'SvcFoo': self.mock_svc1, }) ] assert self.mock_svc1.mock_calls == [ call.get_limits() ] assert self.mock_svc2.mock_calls == [] def test_get_limits_service_with_api(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] res = self.cls.get_limits(service='SvcBar') assert res == {'SvcBar': limits['SvcBar']} assert self.mock_ta.mock_calls == [ call.update_limits({ 'SvcBar': self.mock_svc2, }) ] assert self.mock_svc1.mock_calls == [] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.get_limits() ] def test_find_usage(self): self.cls.find_usage() assert self.mock_svc1.mock_calls == [ call.find_usage() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.find_usage() ] assert self.mock_ta.mock_calls == [ call.update_limits({ 'SvcFoo': self.mock_svc1, 'SvcBar': self.mock_svc2, }) ] def test_find_usage_no_ta(self): self.cls.find_usage(use_ta=False) assert self.mock_svc1.mock_calls == [ call.find_usage() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.find_usage() ] assert self.mock_ta.mock_calls == [] def test_find_usage_service(self): self.cls.find_usage(service='SvcFoo') assert self.mock_svc1.mock_calls == [ call.find_usage() ] assert self.mock_svc2.mock_calls == [] assert self.mock_ta.mock_calls == [ call.update_limits({'SvcFoo': self.mock_svc1}) ] def test_find_usage_service_with_api(self): self.cls.find_usage(service='SvcBar') assert self.mock_svc1.mock_calls == [] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.find_usage() ] assert self.mock_ta.mock_calls == [ call.update_limits({'SvcBar': self.mock_svc2}) ] def test_set_threshold_overrides(self): limits = sample_limits() limits['SvcFoo']['zz3'] = AwsLimit( 'zz3', self.mock_svc1, 1, 2, 3, ) self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] overrides = { 'SvcBar': { 'barlimit1': { 'warning': { 'percent': 10, 'count': 12 }, 'critical': { 'percent': 14, 'count': 16 } }, 'bar limit2': { 'critical': { 'count': 15, } }, 'zz3': { 'warning': { 'count': 41 }, 'critical': { 'percent': 52 } } }, 'SvcFoo': { 'foo limit3': { 'warning': { 'percent': 91 }, } }, } self.cls.set_threshold_overrides(overrides) assert self.mock_svc1.mock_calls == [ call.set_threshold_override( 'foo limit3', warn_percent=91, ) ] assert self.mock_svc2.mock_calls == [ call.set_threshold_override( 'bar limit2', crit_count=15 ), call.set_threshold_override( 'barlimit1', warn_percent=10, warn_count=12, crit_percent=14, crit_count=16 ), call.set_threshold_override( 'zz3', warn_count=41, crit_percent=52 ), ] def test_set_limit_overrides(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] overrides = { 'SvcBar': { 'barlimit1': 100, }, 'SvcFoo': { 'foo limit3': 99, }, } self.cls.set_limit_overrides(overrides) assert self.mock_svc1.mock_calls == [ call.set_limit_override('foo limit3', 99, override_ta=True) ] assert self.mock_svc2.mock_calls == [ call.set_limit_override('barlimit1', 100, override_ta=True) ] def test_set_limit_overrides_ta(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] overrides = { 'SvcBar': { 'bar limit2': 100, }, 'SvcFoo': { 'foo limit3': 3, }, } self.cls.set_limit_overrides(overrides, override_ta=False) assert self.mock_svc1.mock_calls == [ call.set_limit_override('foo limit3', 3, override_ta=False) ] assert self.mock_svc2.mock_calls == [ call.set_limit_override('bar limit2', 100, override_ta=False) ] def test_set_limit_override(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.cls.set_limit_override('SvcFoo', 'foo limit3', 99) assert self.mock_svc1.mock_calls == [ call.set_limit_override('foo limit3', 99, override_ta=True) ] def test_set_limit_override_ta(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.cls.set_limit_override( 'SvcFoo', 'foo limit3', 99, override_ta=False ) assert self.mock_svc1.mock_calls == [ call.set_limit_override( 'foo limit3', 99, override_ta=False ) ] def test_set_threshold_override(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.cls.set_threshold_override( 'SvcFoo', 'foo limit3', warn_percent=10, warn_count=12, crit_percent=14, crit_count=16 ) assert self.mock_svc1.mock_calls == [ call.set_threshold_override( 'foo limit3', warn_percent=10, warn_count=12, crit_percent=14, crit_count=16 ) ] def test_get_required_iam_policy(self): expected = { 'Version': '2012-10-17', 'Statement': [{ 'Effect': 'Allow', 'Resource': '*', 'Action': [ 'ec2:bar', 'ec2:foo', 'foo:perm1', 'foo:perm2', 'support:*', 'trustedadvisor:Describe*' ], }], } self.mock_svc1.required_iam_permissions.return_value = [ 'ec2:foo', 'ec2:bar', ] self.mock_svc2.required_iam_permissions.return_value = [ 'foo:perm1', 'foo:perm2', ] res = self.cls.get_required_iam_policy() assert res == expected assert self.mock_svc1.mock_calls == [call.required_iam_permissions()] assert self.mock_svc2.mock_calls == [call.required_iam_permissions()] def test_check_thresholds(self): self.mock_svc1.check_thresholds.return_value = { 'foo': 'bar', 'baz': 'blam', } self.mock_svc2.check_thresholds.return_value = {} res = self.cls.check_thresholds() assert res == { 'SvcFoo': { 'foo': 'bar', 'baz': 'blam', } } assert self.mock_ta.mock_calls == [ call.update_limits({ 'SvcFoo': self.mock_svc1, 'SvcBar': self.mock_svc2 }), ] assert self.mock_svc1.mock_calls == [ call.check_thresholds() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.check_thresholds() ] def test_check_thresholds_service(self): self.mock_svc1.check_thresholds.return_value = {'foo': 'bar'} self.mock_svc2.check_thresholds.return_value = {'baz': 'blam'} res = self.cls.check_thresholds(service='SvcFoo') assert res == { 'SvcFoo': { 'foo': 'bar', } } assert self.mock_ta.mock_calls == [ call.update_limits({'SvcFoo': self.mock_svc1}) ] assert self.mock_svc1.mock_calls == [ call.check_thresholds() ] assert self.mock_svc2.mock_calls == [] def test_check_thresholds_service_api(self): self.mock_svc1.check_thresholds.return_value = {'foo': 'bar'} self.mock_svc2.check_thresholds.return_value = {'baz': 'blam'} res = self.cls.check_thresholds(service='SvcBar') assert res == { 'SvcBar': { 'baz': 'blam', } } assert self.mock_ta.mock_calls == [ call.update_limits({'SvcBar': self.mock_svc2}) ] assert self.mock_svc1.mock_calls == [] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.check_thresholds() ] def test_check_thresholds_no_ta(self): self.mock_svc1.check_thresholds.return_value = { 'foo': 'bar', 'baz': 'blam', } self.mock_svc2.check_thresholds.return_value = {} self.cls.use_ta = False res = self.cls.check_thresholds(use_ta=False) assert res == { 'SvcFoo': { 'foo': 'bar', 'baz': 'blam', } } assert self.mock_ta.mock_calls == [] assert self.mock_svc1.mock_calls == [ call.check_thresholds() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.check_thresholds() ]
class TestAwsLimitChecker(object): def setup(self): self.mock_ver_info = Mock( release='1.2.3', url='http://myurl', commit='abcd', tag='mytag', version_str='1.2.3@mytag' ) self.mock_svc1 = Mock(spec_set=_AwsService) self.mock_svc2 = Mock(spec_set=ApiServiceSpec) self.mock_foo = Mock(spec_set=_AwsService) self.mock_bar = Mock(spec_set=_AwsService) self.mock_ta = Mock(spec_set=TrustedAdvisor) self.mock_foo.return_value = self.mock_svc1 self.mock_bar.return_value = self.mock_svc2 self.svcs = {'SvcFoo': self.mock_foo, 'SvcBar': self.mock_bar} with patch.dict('awslimitchecker.checker._services', values=self.svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, autospec=True, ) as mocks: self.mock_logger = mocks['logger'] self.mock_version = mocks['_get_version_info'] self.mock_ta_constr = mocks['TrustedAdvisor'] mocks['TrustedAdvisor'].return_value = self.mock_ta self.mock_version.return_value = self.mock_ver_info self.cls = AwsLimitChecker() def test_init(self): # dict should be of _AwsService instances assert self.cls.services == { 'SvcFoo': self.mock_svc1, 'SvcBar': self.mock_svc2 } # _AwsService instances should exist, but have no other calls assert self.mock_foo.mock_calls == [ call(80, 99, None, None, None, None, None, None) ] assert self.mock_bar.mock_calls == [ call(80, 99, None, None, None, None, None, None) ] assert self.mock_ta_constr.mock_calls == [ call(account_id=None, account_role=None, region=None, external_id=None, mfa_serial_number=None, mfa_token=None) ] assert self.mock_svc1.mock_calls == [] assert self.mock_svc2.mock_calls == [] assert self.cls.ta == self.mock_ta assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info assert self.mock_logger.mock_calls == [] def test_init_AGPL_message(self, capsys): # get rid of the class self.cls = None # clear out/err out, err = capsys.readouterr() # run setup again self.setup() # check out/err out, err = capsys.readouterr() assert out == '' assert (err) == ( "awslimitchecker 1.2.3@mytag is AGPL-licensed free software; " "all users have a right to the full source code of " "this version. See <http://myurl>\n") def test_init_thresholds(self): mock_svc1 = Mock(spec_set=_AwsService) mock_svc2 = Mock(spec_set=_AwsService) mock_foo = Mock(spec_set=_AwsService) mock_bar = Mock(spec_set=_AwsService) mock_ta = Mock(spec_set=TrustedAdvisor) mock_foo.return_value = mock_svc1 mock_bar.return_value = mock_svc2 svcs = {'SvcFoo': mock_foo, 'SvcBar': mock_bar} with patch.dict('awslimitchecker.checker._services', values=svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, autospec=True, ) as mocks: mock_version = mocks['_get_version_info'] mock_version.return_value = self.mock_ver_info mock_ta_constr = mocks['TrustedAdvisor'] mocks['TrustedAdvisor'].return_value = mock_ta cls = AwsLimitChecker( warning_threshold=5, critical_threshold=22, ) # dict should be of _AwsService instances assert cls.services == { 'SvcFoo': mock_svc1, 'SvcBar': mock_svc2 } # _AwsService instances should exist, but have no other calls assert mock_foo.mock_calls == [ call(5, 22, None, None, None, None, None, None) ] assert mock_bar.mock_calls == [ call(5, 22, None, None, None, None, None, None) ] assert mock_ta_constr.mock_calls == [ call(account_id=None, account_role=None, region=None, external_id=None, mfa_serial_number=None, mfa_token=None) ] assert mock_svc1.mock_calls == [] assert mock_svc2.mock_calls == [] assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info def test_init_region(self): mock_svc1 = Mock(spec_set=_AwsService) mock_svc2 = Mock(spec_set=_AwsService) mock_foo = Mock(spec_set=_AwsService) mock_bar = Mock(spec_set=_AwsService) mock_ta = Mock(spec_set=TrustedAdvisor) mock_foo.return_value = mock_svc1 mock_bar.return_value = mock_svc2 svcs = {'SvcFoo': mock_foo, 'SvcBar': mock_bar} with patch.dict('awslimitchecker.checker._services', values=svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, autospec=True, ) as mocks: mock_version = mocks['_get_version_info'] mock_version.return_value = self.mock_ver_info mock_ta_constr = mocks['TrustedAdvisor'] mocks['TrustedAdvisor'].return_value = mock_ta cls = AwsLimitChecker(region='myregion') # dict should be of _AwsService instances assert cls.services == { 'SvcFoo': mock_svc1, 'SvcBar': mock_svc2 } # _AwsService instances should exist, but have no other calls assert mock_foo.mock_calls == [ call(80, 99, None, None, 'myregion', None, None, None) ] assert mock_bar.mock_calls == [ call(80, 99, None, None, 'myregion', None, None, None) ] assert mock_ta_constr.mock_calls == [ call(account_id=None, account_role=None, region='myregion', external_id=None, mfa_serial_number=None, mfa_token=None) ] assert mock_svc1.mock_calls == [] assert mock_svc2.mock_calls == [] assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info def test_init_sts(self): mock_svc1 = Mock(spec_set=_AwsService) mock_svc2 = Mock(spec_set=_AwsService) mock_foo = Mock(spec_set=_AwsService) mock_bar = Mock(spec_set=_AwsService) mock_ta = Mock(spec_set=TrustedAdvisor) mock_foo.return_value = mock_svc1 mock_bar.return_value = mock_svc2 svcs = {'SvcFoo': mock_foo, 'SvcBar': mock_bar} with patch.dict('awslimitchecker.checker._services', values=svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, autospec=True, ) as mocks: mock_version = mocks['_get_version_info'] mock_version.return_value = self.mock_ver_info mock_ta_constr = mocks['TrustedAdvisor'] mocks['TrustedAdvisor'].return_value = mock_ta cls = AwsLimitChecker( account_id='123456789012', account_role='myrole', region='myregion' ) # dict should be of _AwsService instances assert cls.services == { 'SvcFoo': mock_svc1, 'SvcBar': mock_svc2 } # _AwsService instances should exist, but have no other calls assert mock_foo.mock_calls == [ call(80, 99, '123456789012', 'myrole', 'myregion', None, None, None) ] assert mock_bar.mock_calls == [ call(80, 99, '123456789012', 'myrole', 'myregion', None, None, None) ] assert mock_ta_constr.mock_calls == [ call( account_id='123456789012', account_role='myrole', region='myregion', external_id=None, mfa_serial_number=None, mfa_token=None ) ] assert mock_svc1.mock_calls == [] assert mock_svc2.mock_calls == [] assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info def test_init_sts_external_id(self): mock_svc1 = Mock(spec_set=_AwsService) mock_svc2 = Mock(spec_set=_AwsService) mock_foo = Mock(spec_set=_AwsService) mock_bar = Mock(spec_set=_AwsService) mock_ta = Mock(spec_set=TrustedAdvisor) mock_foo.return_value = mock_svc1 mock_bar.return_value = mock_svc2 svcs = {'SvcFoo': mock_foo, 'SvcBar': mock_bar} with patch.dict('awslimitchecker.checker._services', values=svcs, clear=True): with patch.multiple( 'awslimitchecker.checker', logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, autospec=True, ) as mocks: mock_version = mocks['_get_version_info'] mock_version.return_value = self.mock_ver_info mock_ta_constr = mocks['TrustedAdvisor'] mocks['TrustedAdvisor'].return_value = mock_ta cls = AwsLimitChecker( account_id='123456789012', account_role='myrole', region='myregion', external_id='myextid', mfa_serial_number=None, mfa_token=None ) # dict should be of _AwsService instances assert cls.services == { 'SvcFoo': mock_svc1, 'SvcBar': mock_svc2 } # _AwsService instances should exist, but have no other calls assert mock_foo.mock_calls == [ call(80, 99, '123456789012', 'myrole', 'myregion', 'myextid', None, None) ] assert mock_bar.mock_calls == [ call(80, 99, '123456789012', 'myrole', 'myregion', 'myextid', None, None) ] assert mock_ta_constr.mock_calls == [ call( account_id='123456789012', account_role='myrole', region='myregion', external_id='myextid', mfa_serial_number=None, mfa_token=None ) ] assert mock_svc1.mock_calls == [] assert mock_svc2.mock_calls == [] assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info def test_get_version(self): with patch('awslimitchecker.checker._get_version_info', spec_set=_get_version_info) as mock_version: self.cls.vinfo = self.mock_ver_info res = self.cls.get_version() assert res == '1.2.3@mytag' assert mock_version.mock_calls == [] def test_get_project_url(self): with patch('awslimitchecker.checker._get_version_info', spec_set=_get_version_info) as mock_version: self.cls.vinfo = self.mock_ver_info res = self.cls.get_project_url() assert res == 'http://myurl' assert mock_version.mock_calls == [] def test_get_service_names(self): res = self.cls.get_service_names() assert res == ['SvcBar', 'SvcFoo'] def test_get_limits(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] res = self.cls.get_limits() assert res == limits assert self.mock_ta.mock_calls == [ call.update_limits({ 'SvcFoo': self.mock_svc1, 'SvcBar': self.mock_svc2, }) ] assert self.mock_svc1.mock_calls == [ call.get_limits() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.get_limits() ] def test_get_limits_no_ta(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] res = self.cls.get_limits(use_ta=False) assert res == limits assert self.mock_ta.mock_calls == [] assert self.mock_svc1.mock_calls == [ call.get_limits() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.get_limits() ] def test_get_limits_service(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] res = self.cls.get_limits(service='SvcFoo') assert res == {'SvcFoo': limits['SvcFoo']} assert self.mock_ta.mock_calls == [ call.update_limits({ 'SvcFoo': self.mock_svc1, }) ] assert self.mock_svc1.mock_calls == [ call.get_limits() ] assert self.mock_svc2.mock_calls == [] def test_get_limits_service_with_api(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] res = self.cls.get_limits(service='SvcBar') assert res == {'SvcBar': limits['SvcBar']} assert self.mock_ta.mock_calls == [ call.update_limits({ 'SvcBar': self.mock_svc2, }) ] assert self.mock_svc1.mock_calls == [] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.get_limits() ] def test_find_usage(self): self.cls.find_usage() assert self.mock_svc1.mock_calls == [ call.find_usage() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.find_usage() ] assert self.mock_ta.mock_calls == [ call.update_limits({ 'SvcFoo': self.mock_svc1, 'SvcBar': self.mock_svc2, }) ] def test_find_usage_no_ta(self): self.cls.find_usage(use_ta=False) assert self.mock_svc1.mock_calls == [ call.find_usage() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.find_usage() ] assert self.mock_ta.mock_calls == [] def test_find_usage_service(self): self.cls.find_usage(service='SvcFoo') assert self.mock_svc1.mock_calls == [ call.find_usage() ] assert self.mock_svc2.mock_calls == [] assert self.mock_ta.mock_calls == [ call.update_limits({'SvcFoo': self.mock_svc1}) ] def test_find_usage_service_with_api(self): self.cls.find_usage(service='SvcBar') assert self.mock_svc1.mock_calls == [] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.find_usage() ] assert self.mock_ta.mock_calls == [ call.update_limits({'SvcBar': self.mock_svc2}) ] def test_set_threshold_overrides(self): limits = sample_limits() limits['SvcFoo']['zz3'] = AwsLimit( 'zz3', self.mock_svc1, 1, 2, 3, ) self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] overrides = { 'SvcBar': { 'barlimit1': { 'warning': { 'percent': 10, 'count': 12 }, 'critical': { 'percent': 14, 'count': 16 } }, 'bar limit2': { 'critical': { 'count': 15, } }, 'zz3': { 'warning': { 'count': 41 }, 'critical': { 'percent': 52 } } }, 'SvcFoo': { 'foo limit3': { 'warning': { 'percent': 91 }, } }, } self.cls.set_threshold_overrides(overrides) assert self.mock_svc1.mock_calls == [ call.set_threshold_override( 'foo limit3', warn_percent=91, ) ] assert self.mock_svc2.mock_calls == [ call.set_threshold_override( 'bar limit2', crit_count=15 ), call.set_threshold_override( 'barlimit1', warn_percent=10, warn_count=12, crit_percent=14, crit_count=16 ), call.set_threshold_override( 'zz3', warn_count=41, crit_percent=52 ), ] def test_set_limit_overrides(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] overrides = { 'SvcBar': { 'barlimit1': 100, }, 'SvcFoo': { 'foo limit3': 99, }, } self.cls.set_limit_overrides(overrides) assert self.mock_svc1.mock_calls == [ call.set_limit_override('foo limit3', 99, override_ta=True) ] assert self.mock_svc2.mock_calls == [ call.set_limit_override('barlimit1', 100, override_ta=True) ] def test_set_limit_overrides_ta(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.mock_svc2.get_limits.return_value = limits['SvcBar'] overrides = { 'SvcBar': { 'bar limit2': 100, }, 'SvcFoo': { 'foo limit3': 3, }, } self.cls.set_limit_overrides(overrides, override_ta=False) assert self.mock_svc1.mock_calls == [ call.set_limit_override('foo limit3', 3, override_ta=False) ] assert self.mock_svc2.mock_calls == [ call.set_limit_override('bar limit2', 100, override_ta=False) ] def test_set_limit_override(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.cls.set_limit_override('SvcFoo', 'foo limit3', 99) assert self.mock_svc1.mock_calls == [ call.set_limit_override('foo limit3', 99, override_ta=True) ] def test_set_limit_override_ta(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.cls.set_limit_override( 'SvcFoo', 'foo limit3', 99, override_ta=False ) assert self.mock_svc1.mock_calls == [ call.set_limit_override( 'foo limit3', 99, override_ta=False ) ] def test_set_threshold_override(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits['SvcFoo'] self.cls.set_threshold_override( 'SvcFoo', 'foo limit3', warn_percent=10, warn_count=12, crit_percent=14, crit_count=16 ) assert self.mock_svc1.mock_calls == [ call.set_threshold_override( 'foo limit3', warn_percent=10, warn_count=12, crit_percent=14, crit_count=16 ) ] def test_get_required_iam_policy(self): expected = { 'Version': '2012-10-17', 'Statement': [{ 'Effect': 'Allow', 'Resource': '*', 'Action': [ 'ec2:bar', 'ec2:foo', 'foo:perm1', 'foo:perm2', 'support:*', 'trustedadvisor:Describe*' ], }], } self.mock_svc1.required_iam_permissions.return_value = [ 'ec2:foo', 'ec2:bar', ] self.mock_svc2.required_iam_permissions.return_value = [ 'foo:perm1', 'foo:perm2', ] res = self.cls.get_required_iam_policy() assert res == expected assert self.mock_svc1.mock_calls == [call.required_iam_permissions()] assert self.mock_svc2.mock_calls == [call.required_iam_permissions()] def test_check_thresholds(self): self.mock_svc1.check_thresholds.return_value = { 'foo': 'bar', 'baz': 'blam', } self.mock_svc2.check_thresholds.return_value = {} res = self.cls.check_thresholds() assert res == { 'SvcFoo': { 'foo': 'bar', 'baz': 'blam', } } assert self.mock_ta.mock_calls == [ call.update_limits({ 'SvcFoo': self.mock_svc1, 'SvcBar': self.mock_svc2 }), ] assert self.mock_svc1.mock_calls == [ call.check_thresholds() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.check_thresholds() ] def test_check_thresholds_service(self): self.mock_svc1.check_thresholds.return_value = {'foo': 'bar'} self.mock_svc2.check_thresholds.return_value = {'baz': 'blam'} res = self.cls.check_thresholds(service='SvcFoo') assert res == { 'SvcFoo': { 'foo': 'bar', } } assert self.mock_ta.mock_calls == [ call.update_limits({'SvcFoo': self.mock_svc1}) ] assert self.mock_svc1.mock_calls == [ call.check_thresholds() ] assert self.mock_svc2.mock_calls == [] def test_check_thresholds_service_api(self): self.mock_svc1.check_thresholds.return_value = {'foo': 'bar'} self.mock_svc2.check_thresholds.return_value = {'baz': 'blam'} res = self.cls.check_thresholds(service='SvcBar') assert res == { 'SvcBar': { 'baz': 'blam', } } assert self.mock_ta.mock_calls == [ call.update_limits({'SvcBar': self.mock_svc2}) ] assert self.mock_svc1.mock_calls == [] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.check_thresholds() ] def test_check_thresholds_no_ta(self): self.mock_svc1.check_thresholds.return_value = { 'foo': 'bar', 'baz': 'blam', } self.mock_svc2.check_thresholds.return_value = {} self.cls.use_ta = False res = self.cls.check_thresholds(use_ta=False) assert res == { 'SvcFoo': { 'foo': 'bar', 'baz': 'blam', } } assert self.mock_ta.mock_calls == [] assert self.mock_svc1.mock_calls == [ call.check_thresholds() ] assert self.mock_svc2.mock_calls == [ call._update_limits_from_api(), call.check_thresholds() ]
class TestAwsLimitChecker(object): def setup(self): self.mock_ver_info = Mock( release="1.2.3", url="http://myurl", commit="abcd", tag="mytag", version_str="1.2.3@mytag" ) self.mock_svc1 = Mock(spec_set=_AwsService) self.mock_svc2 = Mock(spec_set=_AwsService) self.mock_foo = Mock(spec_set=_AwsService) self.mock_bar = Mock(spec_set=_AwsService) self.mock_ta = Mock(spec_set=TrustedAdvisor) self.mock_foo.return_value = self.mock_svc1 self.mock_bar.return_value = self.mock_svc2 self.svcs = {"SvcFoo": self.mock_foo, "SvcBar": self.mock_bar} with patch.dict("awslimitchecker.checker._services", values=self.svcs, clear=True): with patch.multiple( "awslimitchecker.checker", logger=DEFAULT, _get_version_info=DEFAULT, TrustedAdvisor=DEFAULT, autospec=True, ) as mocks: self.mock_logger = mocks["logger"] self.mock_version = mocks["_get_version_info"] mocks["TrustedAdvisor"].return_value = self.mock_ta self.mock_version.return_value = self.mock_ver_info self.cls = AwsLimitChecker() def test_init(self): # dict should be of _AwsService instances assert self.cls.services == {"SvcFoo": self.mock_svc1, "SvcBar": self.mock_svc2} # _AwsService instances should exist, but have no other calls assert self.mock_foo.mock_calls == [call(80, 99)] assert self.mock_bar.mock_calls == [call(80, 99)] assert self.mock_svc1.mock_calls == [] assert self.mock_svc2.mock_calls == [] assert self.cls.ta == self.mock_ta assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info def test_init_thresholds(self): mock_svc1 = Mock(spec_set=_AwsService) mock_svc2 = Mock(spec_set=_AwsService) mock_foo = Mock(spec_set=_AwsService) mock_bar = Mock(spec_set=_AwsService) mock_foo.return_value = mock_svc1 mock_bar.return_value = mock_svc2 svcs = {"SvcFoo": mock_foo, "SvcBar": mock_bar} with patch.dict("awslimitchecker.checker._services", values=svcs, clear=True): with patch.multiple( "awslimitchecker.checker", logger=DEFAULT, _get_version_info=DEFAULT, autospec=True ) as mocks: mock_version = mocks["_get_version_info"] mock_version.return_value = self.mock_ver_info cls = AwsLimitChecker(warning_threshold=5, critical_threshold=22) # dict should be of _AwsService instances assert cls.services == {"SvcFoo": mock_svc1, "SvcBar": mock_svc2} # _AwsService instances should exist, but have no other calls assert mock_foo.mock_calls == [call(5, 22)] assert mock_bar.mock_calls == [call(5, 22)] assert mock_svc1.mock_calls == [] assert mock_svc2.mock_calls == [] assert self.mock_version.mock_calls == [call()] assert self.cls.vinfo == self.mock_ver_info def test_init_logger(self): """ensure we log a license message""" assert self.mock_logger.mock_calls == [ call.warning( "awslimitchecker %s is AGPL-licensed free software; " "all users have a right to the full source code of " "this version. See <%s>", "1.2.3@mytag", "http://myurl", ) ] def test_get_version(self): with patch("awslimitchecker.checker._get_version_info", spec_set=_get_version_info) as mock_version: self.cls.vinfo = self.mock_ver_info res = self.cls.get_version() assert res == "1.2.3@mytag" assert mock_version.mock_calls == [] def test_get_project_url(self): with patch("awslimitchecker.checker._get_version_info", spec_set=_get_version_info) as mock_version: self.cls.vinfo = self.mock_ver_info res = self.cls.get_project_url() assert res == "http://myurl" assert mock_version.mock_calls == [] def test_get_service_names(self): res = self.cls.get_service_names() assert res == ["SvcBar", "SvcFoo"] def test_get_limits(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits["SvcFoo"] self.mock_svc2.get_limits.return_value = limits["SvcBar"] res = self.cls.get_limits() assert res == limits assert self.mock_ta.mock_calls == [call.update_limits({"SvcFoo": self.mock_svc1, "SvcBar": self.mock_svc2})] def test_get_limits_no_ta(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits["SvcFoo"] self.mock_svc2.get_limits.return_value = limits["SvcBar"] res = self.cls.get_limits(use_ta=False) assert res == limits assert self.mock_ta.mock_calls == [] def test_get_limits_service(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits["SvcFoo"] self.mock_svc2.get_limits.return_value = limits["SvcBar"] res = self.cls.get_limits(service="SvcFoo") assert res == {"SvcFoo": limits["SvcFoo"]} assert self.mock_ta.mock_calls == [call.update_limits({"SvcFoo": self.mock_svc1})] def test_find_usage(self): self.cls.find_usage() assert self.mock_svc1.mock_calls == [call.find_usage()] assert self.mock_svc2.mock_calls == [call.find_usage()] assert self.mock_ta.mock_calls == [call.update_limits({"SvcFoo": self.mock_svc1, "SvcBar": self.mock_svc2})] def test_find_usage_no_ta(self): self.cls.find_usage(use_ta=False) assert self.mock_svc1.mock_calls == [call.find_usage()] assert self.mock_svc2.mock_calls == [call.find_usage()] assert self.mock_ta.mock_calls == [] def test_find_usage_service(self): self.cls.find_usage(service="SvcFoo") assert self.mock_svc1.mock_calls == [call.find_usage()] assert self.mock_svc2.mock_calls == [] assert self.mock_ta.mock_calls == [call.update_limits({"SvcFoo": self.mock_svc1})] def test_set_threshold_overrides(self): limits = sample_limits() limits["SvcFoo"]["zz3"] = AwsLimit("zz3", self.mock_svc1, 1, 2, 3) self.mock_svc1.get_limits.return_value = limits["SvcFoo"] self.mock_svc2.get_limits.return_value = limits["SvcBar"] overrides = { "SvcBar": { "barlimit1": {"warning": {"percent": 10, "count": 12}, "critical": {"percent": 14, "count": 16}}, "bar limit2": {"critical": {"count": 15}}, "zz3": {"warning": {"count": 41}, "critical": {"percent": 52}}, }, "SvcFoo": {"foo limit3": {"warning": {"percent": 91}}}, } self.cls.set_threshold_overrides(overrides) assert self.mock_svc1.mock_calls == [call.set_threshold_override("foo limit3", warn_percent=91)] assert self.mock_svc2.mock_calls == [ call.set_threshold_override("bar limit2", crit_count=15), call.set_threshold_override("barlimit1", warn_percent=10, warn_count=12, crit_percent=14, crit_count=16), call.set_threshold_override("zz3", warn_count=41, crit_percent=52), ] def test_set_limit_overrides(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits["SvcFoo"] self.mock_svc2.get_limits.return_value = limits["SvcBar"] overrides = {"SvcBar": {"barlimit1": 100}, "SvcFoo": {"foo limit3": 99}} self.cls.set_limit_overrides(overrides) assert self.mock_svc1.mock_calls == [call.set_limit_override("foo limit3", 99, override_ta=True)] assert self.mock_svc2.mock_calls == [call.set_limit_override("barlimit1", 100, override_ta=True)] def test_set_limit_overrides_ta(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits["SvcFoo"] self.mock_svc2.get_limits.return_value = limits["SvcBar"] overrides = {"SvcBar": {"bar limit2": 100}, "SvcFoo": {"foo limit3": 3}} self.cls.set_limit_overrides(overrides, override_ta=False) assert self.mock_svc1.mock_calls == [call.set_limit_override("foo limit3", 3, override_ta=False)] assert self.mock_svc2.mock_calls == [call.set_limit_override("bar limit2", 100, override_ta=False)] def test_set_limit_override(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits["SvcFoo"] self.cls.set_limit_override("SvcFoo", "foo limit3", 99) assert self.mock_svc1.mock_calls == [call.set_limit_override("foo limit3", 99, override_ta=True)] def test_set_limit_override_ta(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits["SvcFoo"] self.cls.set_limit_override("SvcFoo", "foo limit3", 99, override_ta=False) assert self.mock_svc1.mock_calls == [call.set_limit_override("foo limit3", 99, override_ta=False)] def test_set_threshold_override(self): limits = sample_limits() self.mock_svc1.get_limits.return_value = limits["SvcFoo"] self.cls.set_threshold_override( "SvcFoo", "foo limit3", warn_percent=10, warn_count=12, crit_percent=14, crit_count=16 ) assert self.mock_svc1.mock_calls == [ call.set_threshold_override("foo limit3", warn_percent=10, warn_count=12, crit_percent=14, crit_count=16) ] def test_get_required_iam_policy(self): expected = { "Version": "2012-10-17", "Statement": [ { "Effect": "Allow", "Resource": "*", "Action": ["ec2:bar", "ec2:foo", "foo:perm1", "foo:perm2", "support:*", "trustedadvisor:Describe*"], } ], } self.mock_svc1.required_iam_permissions.return_value = ["ec2:foo", "ec2:bar"] self.mock_svc2.required_iam_permissions.return_value = ["foo:perm1", "foo:perm2"] res = self.cls.get_required_iam_policy() assert res == expected assert self.mock_svc1.mock_calls == [call.required_iam_permissions()] assert self.mock_svc2.mock_calls == [call.required_iam_permissions()] def test_check_thresholds(self): self.mock_svc1.check_thresholds.return_value = {"foo": "bar", "baz": "blam"} self.mock_svc2.check_thresholds.return_value = {} res = self.cls.check_thresholds() assert res == {"SvcFoo": {"foo": "bar", "baz": "blam"}} assert self.mock_ta.mock_calls == [call.update_limits({"SvcFoo": self.mock_svc1, "SvcBar": self.mock_svc2})] def test_check_thresholds_service(self): self.mock_svc1.check_thresholds.return_value = {"foo": "bar"} self.mock_svc2.check_thresholds.return_value = {"baz": "blam"} res = self.cls.check_thresholds(service="SvcFoo") assert res == {"SvcFoo": {"foo": "bar"}} assert self.mock_ta.mock_calls == [call.update_limits({"SvcFoo": self.mock_svc1})] def test_check_thresholds_no_ta(self): self.mock_svc1.check_thresholds.return_value = {"foo": "bar", "baz": "blam"} self.mock_svc2.check_thresholds.return_value = {} self.cls.use_ta = False res = self.cls.check_thresholds(use_ta=False) assert res == {"SvcFoo": {"foo": "bar", "baz": "blam"}} assert self.mock_ta.mock_calls == []
def check_limits(self, verbose=False): """ Run the actual usage and limit check, with overrides. see: http://awslimitchecker.readthedocs.org/en/latest/python_usage.html#ci-deployment-checks """ # instantiate the class checker = AwsLimitChecker() # set your overrides checker.set_threshold_overrides(AWS_THRESHOLD_OVERRIDES) checker.set_limit_overrides(AWS_LIMIT_OVERRIDES) print("Checking AWS resource usage; WARNING threshold {w}% of " "limit, CRITICAL threshold {c}% of limit".format( w=checker.warning_threshold, c=checker.critical_threshold)) # check usage against thresholds # if we didn't support verbose output, we could just iterate the return # value of this to be a bit more efficient. checker.check_thresholds() # save state for exit code and summary warnings = [] criticals = [] # iterate the results for service, svc_limits in sorted(checker.get_limits().items()): for limit_name, limit in sorted(svc_limits.items()): have_alarms = False # check warnings and criticals for each Limit for warn in limit.get_warnings(): warnings.append(colored("{service} '{limit_name}' usage " "({u}) exceeds warning threshold " "(limit={l})".format( service=service, limit_name=limit_name, u=str(warn), l=limit.get_limit(), ), 'yellow')) have_alarms = True for crit in limit.get_criticals(): criticals.append(colored("{service} '{limit_name}' usage " "({u}) exceeds critical threshold" " (limit={l})".format( service=service, limit_name=limit_name, u=str(crit), l=limit.get_limit(), ), 'red')) have_alarms = True if not have_alarms and verbose: print("{service} '{limit_name}' OK: {u} (limit={l})".format( service=service, limit_name=limit_name, u=limit.get_current_usage_str(), l=limit.get_limit() )) if verbose: print("\n\n") return (warnings, criticals)