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_handle_event_success(self):
with patch.multiple(
pbm,
autospec=True,
logger=DEFAULT,
queues_for_endpoint=DEFAULT,
msg_body_for_event=DEFAULT,
boto3=DEFAULT,
try_enqueue=DEFAULT
) as mocks:
mocks['queues_for_endpoint'].return_value = ['q1']
mocks['msg_body_for_event'].return_value = 'mybody'
mocks['try_enqueue'].return_value = 'msgid'
res = handle_event(self.mock_event, self.mock_context)
assert res == {
'status': 'success',
'message': 'enqueued 1 messages',
'SQSMessageIds': ['msgid']
}
assert mocks['queues_for_endpoint'].mock_calls == [
call(self.mock_event)
]
assert mocks['msg_body_for_event'].mock_calls == [
call(self.mock_event, self.mock_context)
]
assert mocks['try_enqueue'].mock_calls == [
call(mocks['boto3'].client.return_value, 'q1', 'mybody')
]
assert mocks['boto3'].mock_calls == [
call.client('sqs')
]
assert mocks['logger'].mock_calls == []
def test_cloudwatch_connection_needed(self):
mock_conf = Mock(region_name='foo')
mock_cw = Mock(_client_config=mock_conf)
cls = AwsServiceTester(1, 2, {'foo': 'bar'}, None)
assert cls._cloudwatch_client is None
with patch('awslimitchecker.services.base.boto3.client') as m_boto:
m_boto.return_value = mock_cw
res = cls._cloudwatch_connection()
assert res == mock_cw
assert cls._cloudwatch_client == mock_cw
assert m_boto.mock_calls == [call.client('cloudwatch', foo='bar')]
def test_cloudwatch_connection_needed_max_retries(self):
mock_conf = Mock(region_name='foo')
mock_cw = Mock(_client_config=mock_conf)
cls = AwsServiceTester(1, 2, {'foo': 'bar'}, None)
assert cls._cloudwatch_client is None
with patch('awslimitchecker.services.base.boto3.client') as m_boto:
with patch(
'awslimitchecker.connectable.Connectable._max_retries_config',
new_callable=PropertyMock) as m_mrc:
m_mrc.return_value = {'retries': 5}
m_boto.return_value = mock_cw
res = cls._cloudwatch_connection()
assert res == mock_cw
assert cls._cloudwatch_client == mock_cw
assert m_boto.mock_calls == [
call.client('cloudwatch', foo='bar', config={'retries': 5})
]
def test_handle_event(self):
def se_enqueue(conn, qname, msg):
if qname == 'q1':
return 'msgid1'
if qname == 'q2':
raise Exception('foo')
if qname == 'q3':
return 'msgid3'
return 'othermsgid'
with patch.multiple(
pbm,
autospec=True,
logger=DEFAULT,
queues_for_endpoint=DEFAULT,
msg_body_for_event=DEFAULT,
boto3=DEFAULT,
try_enqueue=DEFAULT
) as mocks:
mocks['queues_for_endpoint'].return_value = ['q1', 'q2', 'q3']
mocks['msg_body_for_event'].return_value = 'mybody'
mocks['try_enqueue'].side_effect = se_enqueue
res = handle_event(self.mock_event, self.mock_context)
assert res == {
'status': 'partial',
'message': 'enqueued 2 messages; 1 failed',
'SQSMessageIds': ['msgid1', 'msgid3']
}
assert mocks['queues_for_endpoint'].mock_calls == [
call(self.mock_event)
]
assert mocks['msg_body_for_event'].mock_calls == [
call(self.mock_event, self.mock_context)
]
assert mocks['try_enqueue'].mock_calls == [
call(mocks['boto3'].client.return_value, 'q1', 'mybody'),
call(mocks['boto3'].client.return_value, 'q2', 'mybody'),
call(mocks['boto3'].client.return_value, 'q3', 'mybody'),
]
assert mocks['boto3'].mock_calls == [
call.client('sqs')
]
assert mocks['logger'].mock_calls == [
call.error('Failed enqueueing message in %s:', 'q2', exc_info=1)
]