class Test_TrustedAdvisor(object):

    def setup(self):
        self.mock_conn = Mock()
        self.mock_client_config = Mock()
        type(self.mock_client_config).region_name = 'us-east-1'
        type(self.mock_conn)._client_config = self.mock_client_config
        self.cls = TrustedAdvisor({})
        self.cls.conn = self.mock_conn

        self.mock_svc1 = Mock(spec_set=_AwsService)
        self.mock_svc2 = Mock(spec_set=_AwsService)
        self.services = {
            'SvcFoo': self.mock_svc1,
            'SvcBar': self.mock_svc2,
        }

    def test_init(self):
        cls = TrustedAdvisor({})
        assert cls.conn is None
        assert cls.account_id is None
        assert cls.account_role is None
        assert cls.region == 'us-east-1'
        assert cls.ta_region is None
        assert cls.external_id is None
        assert cls.mfa_serial_number is None
        assert cls.mfa_token is None
        assert cls.all_services == {}
        assert cls.limits_updated is False

    def test_init_sts(self):
        mock_svc = Mock(spec_set=_AwsService)
        mock_svc.get_limits.return_value = {}
        cls = TrustedAdvisor(
            {'foo': mock_svc},
            account_id='aid', account_role='role', region='r'
        )
        assert cls.conn is None
        assert cls.account_id == 'aid'
        assert cls.account_role == 'role'
        assert cls.region == 'us-east-1'
        assert cls.ta_region == 'r'
        assert cls.external_id is None
        assert cls.mfa_serial_number is None
        assert cls.mfa_token is None
        assert cls.all_services == {'foo': mock_svc}
        assert cls.limits_updated is False

    def test_init_sts_external_id(self):
        cls = TrustedAdvisor(
            {}, account_id='aid', account_role='role', region='r',
            external_id='myeid'
        )
        assert cls.conn is None
        assert cls.account_id == 'aid'
        assert cls.account_role == 'role'
        assert cls.region == 'us-east-1'
        assert cls.ta_region == 'r'
        assert cls.external_id == 'myeid'
        assert cls.mfa_serial_number is None
        assert cls.mfa_token is None
        assert cls.limits_updated is False

    def test_update_limits(self):
        mock_results = Mock()
        with patch('%s.connect' % pb, autospec=True) as mock_connect:
            with patch('%s._poll' % pb, autospec=True) as mock_poll:
                with patch('%s._update_services' % pb,
                           autospec=True) as mock_update_services:
                    mock_poll.return_value = mock_results
                    self.cls.update_limits()
        assert mock_connect.mock_calls == [call(self.cls)]
        assert mock_poll.mock_calls == [call(self.cls)]
        assert mock_update_services.mock_calls == [
            call(self.cls, mock_results)
        ]

    def test_update_limits_again(self):
        mock_results = Mock()
        self.cls.limits_updated = True
        with patch('%s.connect' % pb, autospec=True) as mock_connect:
            with patch('%s._poll' % pb, autospec=True) as mock_poll:
                with patch('%s._update_services' % pb,
                           autospec=True) as mock_update_services:
                    with patch('%s.logger' % pbm) as mock_logger:
                        mock_poll.return_value = mock_results
                        self.cls.update_limits()
        assert mock_connect.mock_calls == []
        assert mock_poll.mock_calls == []
        assert mock_update_services.mock_calls == []
        assert mock_logger.mock_calls == [
            call.debug('Already polled TA; skipping update')
        ]

    def test_get_limit_check_id(self):
        api_resp = {
            'checks': [
                {
                    'category': 'performance',
                    'name': 'Service Limits',
                    'id': 'bar',
                    'metadata': [
                        'Region',
                        'Service',
                        'Limit Name',
                        'Limit Amount',
                        'Current Usage',
                        'Status'
                    ],
                },
                {
                    'category': 'fault_tolerance',
                    'name': 'ELB Cross-Zone Load Balancing',
                    'id': 'foo',
                },
            ]
        }
        self.mock_conn.describe_trusted_advisor_checks.return_value = api_resp
        res = self.cls._get_limit_check_id()
        assert res == (
            'bar',
            [
                'Region',
                'Service',
                'Limit Name',
                'Limit Amount',
                'Current Usage',
                'Status'
            ]
        )
        assert self.mock_conn.mock_calls == [
            call.describe_trusted_advisor_checks(language='en')
        ]

    def test_get_limit_check_id_none(self):
        api_resp = {
            'checks': [
                {
                    'category': 'performance',
                    'name': 'Something Else',
                    'id': 'bar',
                },
                {
                    'category': 'fault_tolerance',
                    'name': 'ELB Cross-Zone Load Balancing',
                    'id': 'foo',
                },
            ]
        }
        self.mock_conn.describe_trusted_advisor_checks.return_value = api_resp
        res = self.cls._get_limit_check_id()
        assert res == (None, None)
        assert self.mock_conn.mock_calls == [
            call.describe_trusted_advisor_checks(language='en')
        ]

    def test_get_limit_check_id_subscription_required(self):

        def se_api(language=None):
            response = {
                'ResponseMetadata': {
                    'HTTPStatusCode': 400,
                    'RequestId': '3cc9b2a8-c6e5-11e5-bc1d-b13dcea36176'
                },
                'Error': {
                    'Message': 'AWS Premium Support Subscription is required '
                               'to use this service.',
                    'Code': 'SubscriptionRequiredException'
                }
            }
            raise ClientError(response, 'operation')

        assert self.cls.have_ta is True
        self.mock_conn.describe_trusted_advisor_checks.side_effect = se_api
        with patch('awslimitchecker.trustedadvisor'
                   '.logger', autospec=True) as mock_logger:
            res = self.cls._get_limit_check_id()
        assert self.cls.have_ta is False
        assert res == (None, None)
        assert self.mock_conn.mock_calls == [
            call.describe_trusted_advisor_checks(language='en')
        ]
        assert mock_logger.mock_calls == [
            call.debug("Querying Trusted Advisor checks"),
            call.warning("Cannot check TrustedAdvisor: %s",
                         'AWS Premium Support Subscription is required to '
                         'use this service.')
        ]

    def test_get_limit_check_id_other_exception(self):

        def se_api(language=None):
            response = {
                'ResponseMetadata': {
                    'HTTPStatusCode': 400,
                    'RequestId': '3cc9b2a8-c6e5-11e5-bc1d-b13dcea36176'
                },
                'Error': {
                    'Message': 'foo',
                    'Code': 'SomeOtherException'
                }
            }
            raise ClientError(response, 'operation')

        self.mock_conn.describe_trusted_advisor_checks.side_effect = se_api
        with pytest.raises(ClientError) as excinfo:
            self.cls._get_limit_check_id()
        assert self.mock_conn.mock_calls == [
            call.describe_trusted_advisor_checks(language='en')
        ]
        assert excinfo.value.response['ResponseMetadata'][
                   'HTTPStatusCode'] == 400
        assert excinfo.value.response['Error']['Message'] == 'foo'
        assert excinfo.value.response['Error']['Code'] == 'SomeOtherException'

    def test_poll_id_none(self):
        tmp = self.mock_conn.describe_trusted_advisor_check_result
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            mock_id.return_value = None
            res = self.cls._poll()
        assert tmp.mock_calls == []
        assert res is None

    def test_poll(self):
        tmp = self.mock_conn.describe_trusted_advisor_check_result
        poll_return_val = {
            'result': {
                'timestamp': '2015-06-15T20:27:42Z',
                'flaggedResources': [
                    {
                        'status': 'ok',
                        'resourceId': 'resid1',
                        'isSuppressed': False,
                        'region': 'us-west-2',
                        'metadata': [
                            'us-west-2',
                            'AutoScaling',
                            'Auto Scaling groups',
                            '20',
                            '2',
                            'Green'
                        ]
                    },
                    {
                        'status': 'ok',
                        'resourceId': 'resid2',
                        'isSuppressed': False,
                        'region': 'us-east-1',
                        'metadata': [
                            'us-east-1',
                            'AutoScaling',
                            'Launch configurations',
                            '20',
                            '18',
                            'Yellow'
                        ]
                    },
                    {
                        'status': 'ok',
                        'resourceId': 'resid3',
                        'isSuppressed': False,
                        'region': 'us-east-1',
                        'metadata': [
                            'us-west-2',
                            'AutoScaling',
                            'Auto Scaling groups',
                            '40',
                            '10',
                            'Green'
                        ]
                    },
                    {
                        'status': 'ok',
                        'resourceId': 'resid4',
                        'isSuppressed': False,
                        'metadata': [
                            '-',
                            'IAM',
                            'Users',
                            '5000',
                            '152',
                            'Green'
                        ]
                    },
                ]
            }
        }
        tmp.return_value = poll_return_val
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            mock_id.return_value = (
                'foo',
                [
                    'Region',
                    'Service',
                    'Limit Name',
                    'Limit Amount',
                    'Current Usage',
                    'Status'
                ]
            )
            res = self.cls._poll()
        assert self.mock_conn.mock_calls == [
            call.describe_trusted_advisor_check_result(
                checkId='foo', language='en'
            )
        ]
        assert mock_id.mock_calls == [call(self.cls)]
        assert res == {
            'AutoScaling': {
                'Launch configurations': 20,
                'Auto Scaling groups': 40,
            },
            'IAM': {
                'Users': 5000
            }
        }

    def test_poll_region(self):
        tmp = self.mock_conn.describe_trusted_advisor_check_result
        self.cls.ta_region = 'us-west-2'
        poll_return_value = {
            'result': {
                'timestamp': '2015-06-15T20:27:42Z',
                'flaggedResources': [
                    {
                        'status': 'ok',
                        'resourceId': 'resid1',
                        'isSuppressed': False,
                        'region': 'us-west-2',
                        'metadata': [
                            'us-west-2',
                            'AutoScaling',
                            'Auto Scaling groups',
                            '20',
                            '2',
                            'Green'
                        ]
                    },
                    {
                        'status': 'ok',
                        'resourceId': 'resid2',
                        'isSuppressed': False,
                        'region': 'us-east-1',
                        'metadata': [
                            'us-east-1',
                            'AutoScaling',
                            'Launch configurations',
                            '20',
                            '18',
                            'Yellow'
                        ]
                    },
                    {
                        'status': 'ok',
                        'resourceId': 'resid3',
                        'isSuppressed': False,
                        'region': 'us-east-1',
                        'metadata': [
                            'us-west-2',
                            'AutoScaling',
                            'Auto Scaling groups',
                            '40',
                            '10',
                            'Green'
                        ]
                    },
                    {
                        'status': 'ok',
                        'resourceId': 'resid4',
                        'isSuppressed': False,
                        'metadata': [
                            '-',
                            'IAM',
                            'Users',
                            '5000',
                            '152',
                            'Green'
                        ]
                    },
                ]
            }
        }
        tmp.return_value = poll_return_value
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            mock_id.return_value = (
                'foo',
                [
                    'Region',
                    'Service',
                    'Limit Name',
                    'Limit Amount',
                    'Current Usage',
                    'Status'
                ]
            )
            res = self.cls._poll()
        assert self.mock_conn.mock_calls == [
            call.describe_trusted_advisor_check_result(
                checkId='foo', language='en'
            )
        ]
        assert mock_id.mock_calls == [call(self.cls)]
        assert res == {
            'AutoScaling': {
                'Auto Scaling groups': 20,
            },
            'IAM': {
                'Users': 5000
            }
        }

    def test_poll_dont_have_ta(self):
        self.cls.have_ta = False
        tmp = self.mock_conn.describe_trusted_advisor_check_result
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            with patch('awslimitchecker.trustedadvisor'
                       '.logger', autospec=True) as mock_logger:
                res = self.cls._poll()
        assert self.mock_conn.mock_calls == []
        assert tmp.mock_calls == []
        assert mock_id.mock_calls == [
            call(self.cls)
        ]
        assert mock_logger.mock_calls == [
            call.info('Beginning TrustedAdvisor poll'),
            call.info('TrustedAdvisor.have_ta is False; not polling TA')
        ]
        assert res == {}

    def test_update_services(self):
        mock_as_foo = Mock(spec_set=AwsLimit)
        mock_as_bar = Mock(spec_set=AwsLimit)
        mock_ec2_baz = Mock(spec_set=AwsLimit)
        mock_vpc = Mock(spec_set=AwsLimit)
        ta_services = {
            'AutoScaling': {
                'foo': mock_as_foo,
                'bar': mock_as_bar
            },
            'EC2': {
                'baz': mock_ec2_baz
            },
            'VPC': {
                'VPC Elastic IP addresses (EIPs)': mock_vpc
            },
        }
        ta_results = {
            'AutoScaling': {
                'foo': 20,
                'bar': 40,
            },
            'EC2': {
                'baz': 5,
                'blam': 10,
            },
            'OtherService': {
                'blarg': 1,
            },
            'VPC': {
                'VPC Elastic IP addresses (EIPs)': 11,
            }
        }
        with patch('awslimitchecker.trustedadvisor'
                   '.logger', autospec=True) as mock_logger:
            self.cls.ta_services = ta_services
            self.cls._update_services(ta_results)
        assert mock_logger.mock_calls == [
            call.debug("Updating TA limits on all services"),
            call.info("TrustedAdvisor returned check results for unknown "
                      "limit '%s' (service %s)", 'blam', 'EC2'),
            call.info("TrustedAdvisor returned check results for unknown "
                      "service '%s'", 'OtherService'),
            call.info("Done updating TA limits on all services"),
        ]
        assert mock_as_foo.mock_calls == [
            call._set_ta_limit(20)
        ]
        assert mock_as_bar.mock_calls == [
            call._set_ta_limit(40)
        ]
        assert mock_ec2_baz.mock_calls == [
            call._set_ta_limit(5)
        ]
        assert mock_vpc.mock_calls == [
            call._set_ta_limit(11)
        ]

    def test_make_ta_service_dict(self):
        mock_ec2 = Mock(spec_set=_AwsService)
        mock_el1 = Mock(spec_set=AwsLimit)
        type(mock_el1).name = 'el1'
        type(mock_el1).ta_service_name = 'EC2'
        type(mock_el1).ta_limit_name = 'el1'
        mock_el2 = Mock(spec_set=AwsLimit)
        type(mock_el2).name = 'el2'
        type(mock_el2).ta_service_name = 'Foo'
        type(mock_el2).ta_limit_name = 'el2'
        mock_ec2.get_limits.return_value = {
            'mock_el1': mock_el1,
            'mock_el2': mock_el2
        }

        mock_vpc = Mock(spec_set=_AwsService)
        mock_vl1 = Mock(spec_set=AwsLimit)
        type(mock_vl1).name = 'vl1'
        type(mock_vl1).ta_service_name = 'VPC'
        type(mock_vl1).ta_limit_name = 'other name'
        mock_vl2 = Mock(spec_set=AwsLimit)
        type(mock_vl2).name = 'vl2'
        type(mock_vl2).ta_service_name = 'Foo'
        type(mock_vl2).ta_limit_name = 'other limit'
        mock_vpc.get_limits.return_value = {
            'mock_vl1': mock_vl1,
            'mock_vl2': mock_vl2
        }

        svcs = {
            'EC2': mock_ec2,
            'VPC': mock_vpc
        }

        expected = {
            'EC2': {
                'el1': mock_el1
            },
            'VPC': {
                'other name': mock_vl1
            },
            'Foo': {
                'el2': mock_el2,
                'other limit': mock_vl2
            }
        }
        self.cls.all_services = svcs
        assert self.cls._make_ta_service_dict() == expected
class Test_TrustedAdvisor(object):
    def setup(self):
        self.mock_conn = Mock(spec_set=SupportConnection)
        type(self.mock_conn).region = RegionInfo(name='us-east-1')
        self.cls = TrustedAdvisor()
        self.cls.conn = self.mock_conn

        self.mock_svc1 = Mock(spec_set=_AwsService)
        self.mock_svc2 = Mock(spec_set=_AwsService)
        self.services = {
            'SvcFoo': self.mock_svc1,
            'SvcBar': self.mock_svc2,
        }

    def test_init(self):
        cls = TrustedAdvisor()
        assert cls.conn is None
        assert cls.account_id is None
        assert cls.account_role is None
        assert cls.region == 'us-east-1'
        assert cls.ta_region is None
        assert cls.external_id is None

    def test_init_sts(self):
        cls = TrustedAdvisor(account_id='aid', account_role='role', region='r')
        assert cls.conn is None
        assert cls.account_id == 'aid'
        assert cls.account_role == 'role'
        assert cls.region == 'us-east-1'
        assert cls.ta_region == 'r'
        assert cls.external_id is None

    def test_init_sts_external_id(self):
        cls = TrustedAdvisor(account_id='aid',
                             account_role='role',
                             region='r',
                             external_id='myeid')
        assert cls.conn is None
        assert cls.account_id == 'aid'
        assert cls.account_role == 'role'
        assert cls.region == 'us-east-1'
        assert cls.ta_region == 'r'
        assert cls.external_id == 'myeid'

    def test_connect(self):
        cls = TrustedAdvisor()
        mock_conn = Mock(spec_set=SupportConnection, name='mock_conn')
        with patch('awslimitchecker.trustedadvisor.boto.connect_support'
                   '',
                   autospec=True) as mock_connect:
            mock_connect.return_value = mock_conn
            cls.connect()
        assert cls.conn == mock_conn
        assert mock_connect.mock_calls == [call()]

    def test_connect_region(self):
        cls = TrustedAdvisor(account_id='foo', account_role='bar', region='re')
        mock_conn = Mock(spec_set=SupportConnection, name='mock_conn')
        mock_conn_via = Mock(spec_set=SupportConnection, name='mock_conn')
        with patch('awslimitchecker.trustedadvisor.TrustedAdvisor.connect_via'
                   '') as mock_connect_via:
            mock_connect_via.return_value = mock_conn_via
            with patch(
                    'awslimitchecker.trustedadvisor.boto.connect_support'
                    '',
                    autospec=True) as mock_connect:
                mock_connect.return_value = mock_conn
                cls.connect()
        assert cls.conn == mock_conn_via
        assert mock_connect.mock_calls == []
        assert mock_connect_via.mock_calls == [call(connect_to_region)]

    def test_connect_again(self):
        cls = TrustedAdvisor()
        mock_original_conn = Mock(spec_set=SupportConnection)
        cls.conn = mock_original_conn
        mock_conn = Mock(spec_set=SupportConnection)
        with patch('awslimitchecker.trustedadvisor.boto.connect_support'
                   '') as mock_connect:
            mock_connect.return_value = mock_conn
            cls.connect()
        assert cls.conn == mock_original_conn
        assert mock_connect.mock_calls == []

    def test_update_limits(self):
        mock_results = Mock()
        with patch('%s.connect' % pb, autospec=True) as mock_connect:
            with patch('%s._poll' % pb, autospec=True) as mock_poll:
                with patch('%s._update_services' % pb,
                           autospec=True) as mock_update_services:
                    mock_poll.return_value = mock_results
                    self.cls.update_limits(self.services)
        assert mock_connect.mock_calls == [call(self.cls)]
        assert mock_poll.mock_calls == [call(self.cls)]
        assert mock_update_services.mock_calls == [
            call(self.cls, mock_results, self.services)
        ]

    def test_get_limit_check_id(self):
        api_resp = {
            'checks': [
                {
                    'category':
                    'performance',
                    'name':
                    'Service Limits',
                    'id':
                    'bar',
                    'metadata': [
                        'Region', 'Service', 'Limit Name', 'Limit Amount',
                        'Current Usage', 'Status'
                    ],
                },
                {
                    'category': 'fault_tolerance',
                    'name': 'ELB Cross-Zone Load Balancing',
                    'id': 'foo',
                },
            ]
        }
        with patch('%s.boto_query_wrapper' % pbm) as mock_wrapper:
            mock_wrapper.return_value = api_resp
            res = self.cls._get_limit_check_id()
        assert res == ('bar', [
            'Region', 'Service', 'Limit Name', 'Limit Amount', 'Current Usage',
            'Status'
        ])
        assert self.mock_conn.mock_calls == []
        assert mock_wrapper.mock_calls == [
            call(self.mock_conn.describe_trusted_advisor_checks, 'en')
        ]

    def test_get_limit_check_id_none(self):
        api_resp = {
            'checks': [
                {
                    'category': 'performance',
                    'name': 'Something Else',
                    'id': 'bar',
                },
                {
                    'category': 'fault_tolerance',
                    'name': 'ELB Cross-Zone Load Balancing',
                    'id': 'foo',
                },
            ]
        }
        with patch('%s.boto_query_wrapper' % pbm) as mock_wrapper:
            mock_wrapper.return_value = api_resp
            res = self.cls._get_limit_check_id()
        assert res == (None, None)
        assert self.mock_conn.mock_calls == []
        assert mock_wrapper.mock_calls == [
            call(self.mock_conn.describe_trusted_advisor_checks, 'en')
        ]

    def test_get_limit_check_id_subscription_required(self):
        def se_api(foo, language):
            status = 400
            reason = 'Bad Request'
            body = {
                'message': 'AWS Premium Support Subscription is required to '
                'use this service.',
                '__type': 'SubscriptionRequiredException'
            }
            raise JSONResponseError(status, reason, body)

        assert self.cls.have_ta is True
        with patch('awslimitchecker.trustedadvisor'
                   '.logger', autospec=True) as mock_logger:
            with patch('%s.boto_query_wrapper' % pbm) as mock_wrapper:
                mock_wrapper.side_effect = se_api
                res = self.cls._get_limit_check_id()
        assert self.cls.have_ta is False
        assert res == (None, None)
        assert self.mock_conn.mock_calls == []
        assert mock_wrapper.mock_calls == [
            call(self.mock_conn.describe_trusted_advisor_checks, 'en')
        ]
        assert mock_logger.mock_calls == [
            call.debug("Querying Trusted Advisor checks"),
            call.warning(
                "Cannot check TrustedAdvisor: %s", "AWS Premium Support "
                "Subscription is required to use this service.")
        ]

    def test_get_limit_check_id_other_exception(self):
        def se_api(foo, language):
            status = 400
            reason = 'foobar'
            body = {'message': 'other message', '__type': 'OtherException'}
            raise JSONResponseError(status, reason, body)

        with pytest.raises(BotoServerError) as excinfo:
            with patch('%s.boto_query_wrapper' % pbm) as mock_wrapper:
                mock_wrapper.side_effect = se_api
                self.cls._get_limit_check_id()
        assert self.mock_conn.mock_calls == []
        assert mock_wrapper.mock_calls == [
            call(self.mock_conn.describe_trusted_advisor_checks, 'en')
        ]
        assert excinfo.value.status == 400
        assert excinfo.value.reason == 'foobar'
        assert excinfo.value.body['__type'] == 'OtherException'

    def test_poll_id_none(self):
        tmp = self.mock_conn.describe_trusted_advisor_check_result
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            with patch('%s.boto_query_wrapper' % pbm) as mock_wrapper:
                mock_id.return_value = None
                res = self.cls._poll()
        assert tmp.mock_calls == []
        assert mock_wrapper.mock_calls == []
        assert res is None

    def test_poll(self):
        tmp = self.mock_conn.describe_trusted_advisor_check_result
        poll_return_val = {
            'result': {
                'timestamp':
                '2015-06-15T20:27:42Z',
                'flaggedResources': [
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid1',
                        'isSuppressed':
                        False,
                        'region':
                        'us-west-2',
                        'metadata': [
                            'us-west-2', 'AutoScaling', 'Auto Scaling groups',
                            '20', '2', 'Green'
                        ]
                    },
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid2',
                        'isSuppressed':
                        False,
                        'region':
                        'us-east-1',
                        'metadata': [
                            'us-east-1', 'AutoScaling',
                            'Launch configurations', '20', '18', 'Yellow'
                        ]
                    },
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid3',
                        'isSuppressed':
                        False,
                        'region':
                        'us-east-1',
                        'metadata': [
                            'us-west-2', 'AutoScaling', 'Auto Scaling groups',
                            '40', '10', 'Green'
                        ]
                    },
                ]
            }
        }
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            with patch('%s.boto_query_wrapper' % pbm) as mock_wrapper:
                mock_id.return_value = ('foo', [
                    'Region', 'Service', 'Limit Name', 'Limit Amount',
                    'Current Usage', 'Status'
                ])
                mock_wrapper.return_value = poll_return_val
                res = self.cls._poll()
        assert tmp.mock_calls == []
        assert mock_wrapper.mock_calls == [call(tmp, 'foo')]
        assert mock_id.mock_calls == [call(self.cls)]
        assert res == {
            'AutoScaling': {
                'Launch configurations': 20,
                'Auto Scaling groups': 40,
            }
        }

    def test_poll_region(self):
        tmp = self.mock_conn.describe_trusted_advisor_check_result
        self.cls.ta_region = 'us-west-2'
        poll_return_value = {
            'result': {
                'timestamp':
                '2015-06-15T20:27:42Z',
                'flaggedResources': [
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid1',
                        'isSuppressed':
                        False,
                        'region':
                        'us-west-2',
                        'metadata': [
                            'us-west-2', 'AutoScaling', 'Auto Scaling groups',
                            '20', '2', 'Green'
                        ]
                    },
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid2',
                        'isSuppressed':
                        False,
                        'region':
                        'us-east-1',
                        'metadata': [
                            'us-east-1', 'AutoScaling',
                            'Launch configurations', '20', '18', 'Yellow'
                        ]
                    },
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid3',
                        'isSuppressed':
                        False,
                        'region':
                        'us-east-1',
                        'metadata': [
                            'us-west-2', 'AutoScaling', 'Auto Scaling groups',
                            '40', '10', 'Green'
                        ]
                    },
                ]
            }
        }
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            with patch('%s.boto_query_wrapper' % pbm) as mock_wrapper:
                mock_id.return_value = ('foo', [
                    'Region', 'Service', 'Limit Name', 'Limit Amount',
                    'Current Usage', 'Status'
                ])
                mock_wrapper.return_value = poll_return_value
                res = self.cls._poll()
        assert tmp.mock_calls == []
        assert mock_wrapper.mock_calls == [call(tmp, 'foo')]
        assert mock_id.mock_calls == [call(self.cls)]
        assert res == {
            'AutoScaling': {
                'Auto Scaling groups': 20,
            }
        }

    def test_poll_dont_have_ta(self):
        self.cls.have_ta = False
        tmp = self.mock_conn.describe_trusted_advisor_check_result
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            with patch('%s.boto_query_wrapper' % pbm) as mock_wrapper:
                with patch('awslimitchecker.trustedadvisor'
                           '.logger',
                           autospec=True) as mock_logger:
                    res = self.cls._poll()
        assert mock_wrapper.mock_calls == []
        assert tmp.mock_calls == []
        assert mock_id.mock_calls == [call(self.cls)]
        assert mock_logger.mock_calls == [
            call.info('Beginning TrustedAdvisor poll'),
            call.info('TrustedAdvisor.have_ta is False; not polling TA')
        ]
        assert res == {}

    def test_update_services(self):
        def se_set(lname, val):
            if lname == 'blam':
                raise ValueError("foo")

        mock_autoscale = Mock(spec_set=_AwsService)
        mock_ec2 = Mock(spec_set=_AwsService)
        mock_ec2._set_ta_limit.side_effect = se_set
        mock_vpc = Mock(spec_set=_AwsService)
        services = {
            'AutoScaling': mock_autoscale,
            'EC2': mock_ec2,
            'VPC': mock_vpc,
        }
        ta_results = {
            'AutoScaling': {
                'foo': 20,
                'bar': 40,
            },
            'EC2': {
                'baz': 5,
                'blam': 10,
            },
            'OtherService': {
                'blarg': 1,
            },
            'VPC': {
                'VPC Elastic IP addresses (EIPs)': 11,
            }
        }
        with patch('awslimitchecker.trustedadvisor'
                   '.logger', autospec=True) as mock_logger:
            self.cls._update_services(ta_results, services)
        assert mock_logger.mock_calls == [
            call.debug("Updating TA limits on all services"),
            call.info(
                "TrustedAdvisor returned check results for unknown "
                "limit '%s' (service %s)", 'blam', 'EC2'),
            call.info(
                "TrustedAdvisor returned check results for unknown "
                "service '%s'", 'OtherService'),
            call.info("Done updating TA limits on all services"),
        ]
        assert mock_autoscale.mock_calls == [
            call._set_ta_limit('bar', 40),
            call._set_ta_limit('foo', 20),
        ]
        assert mock_ec2.mock_calls == [
            call._set_ta_limit('baz', 5),
            call._set_ta_limit('blam', 10),
            call._set_ta_limit('VPC Elastic IP addresses (EIPs)', 11)
        ]
class Test_TrustedAdvisor(object):

    def setup(self):
        self.mock_conn = Mock(spec_set=SupportConnection)
        type(self.mock_conn).region = RegionInfo(name='us-east-1')
        self.cls = TrustedAdvisor()
        self.cls.conn = self.mock_conn

        self.mock_svc1 = Mock(spec_set=_AwsService)
        self.mock_svc2 = Mock(spec_set=_AwsService)
        self.services = {
            'SvcFoo': self.mock_svc1,
            'SvcBar': self.mock_svc2,
        }

    def test_init(self):
        cls = TrustedAdvisor()
        assert cls.conn is None

    def test_connect(self):
        cls = TrustedAdvisor()
        mock_conn = Mock(spec_set=SupportConnection, name='mock_conn')
        with patch('awslimitchecker.trustedadvisor.boto.connect_support'
                   '', autospec=True) as mock_connect:
            mock_connect.return_value = mock_conn
            cls.connect()
        assert cls.conn == mock_conn
        assert mock_connect.mock_calls == [call()]

    def test_connect_again(self):
        cls = TrustedAdvisor()
        mock_original_conn = Mock(spec_set=SupportConnection)
        cls.conn = mock_original_conn
        mock_conn = Mock(spec_set=SupportConnection)
        with patch('awslimitchecker.trustedadvisor.boto.connect_support'
                   '') as mock_connect:
            mock_connect.return_value = mock_conn
            cls.connect()
        assert cls.conn == mock_original_conn
        assert mock_connect.mock_calls == []

    def test_update_limits(self):
        mock_results = Mock()
        with patch('%s.connect' % pb, autospec=True) as mock_connect:
            with patch('%s._poll' % pb, autospec=True) as mock_poll:
                with patch('%s._update_services' % pb,
                           autospec=True) as mock_update_services:
                    mock_poll.return_value = mock_results
                    self.cls.update_limits(self.services)
        assert mock_connect.mock_calls == [call(self.cls)]
        assert mock_poll.mock_calls == [call(self.cls)]
        assert mock_update_services.mock_calls == [
            call(self.cls, mock_results, self.services)
        ]

    def test_get_limit_check_id(self):
        api_resp = {
            'checks': [
                {
                    'category': 'performance',
                    'name': 'Service Limits',
                    'id': 'bar',
                    'metadata': [
                        'Region',
                        'Service',
                        'Limit Name',
                        'Limit Amount',
                        'Current Usage',
                        'Status'
                    ],
                },
                {
                    'category': 'fault_tolerance',
                    'name': 'ELB Cross-Zone Load Balancing',
                    'id': 'foo',
                },
            ]
        }
        self.mock_conn.describe_trusted_advisor_checks.return_value = api_resp
        res = self.cls._get_limit_check_id()
        assert res == (
            'bar',
            [
                'Region',
                'Service',
                'Limit Name',
                'Limit Amount',
                'Current Usage',
                'Status'
            ]
        )
        assert self.mock_conn.mock_calls == [
            call.describe_trusted_advisor_checks('en')
        ]

    def test_get_limit_check_id_none(self):
        api_resp = {
            'checks': [
                {
                    'category': 'performance',
                    'name': 'Something Else',
                    'id': 'bar',
                },
                {
                    'category': 'fault_tolerance',
                    'name': 'ELB Cross-Zone Load Balancing',
                    'id': 'foo',
                },
            ]
        }
        self.mock_conn.describe_trusted_advisor_checks.return_value = api_resp
        res = self.cls._get_limit_check_id()
        assert res == (None, None)
        assert self.mock_conn.mock_calls == [
            call.describe_trusted_advisor_checks('en')
        ]

    def test_get_limit_check_id_subscription_required(self):

        def se_api(language):
            status = 400
            reason = 'Bad Request'
            body = {
                'message': 'AWS Premium Support Subscription is required to '
                'use this service.',
                '__type': 'SubscriptionRequiredException'
            }
            raise JSONResponseError(status, reason, body)

        self.mock_conn.describe_trusted_advisor_checks.side_effect = se_api
        assert self.cls.have_ta is True
        with patch('awslimitchecker.trustedadvisor'
                   '.logger', autospec=True) as mock_logger:
            res = self.cls._get_limit_check_id()
        assert self.cls.have_ta is False
        assert res == (None, None)
        assert self.mock_conn.mock_calls == [
            call.describe_trusted_advisor_checks('en')
        ]
        assert mock_logger.mock_calls == [
            call.debug("Querying Trusted Advisor checks"),
            call.warning("Cannot check TrustedAdvisor: %s",
                         "AWS Premium Support "
                         "Subscription is required to use this service.")
        ]

    def test_get_limit_check_id_other_exception(self):

        def se_api(language):
            status = 400
            reason = 'foobar'
            body = {
                'message': 'other message',
                '__type': 'OtherException'
            }
            raise JSONResponseError(status, reason, body)

        self.mock_conn.describe_trusted_advisor_checks.side_effect = se_api
        with pytest.raises(BotoServerError) as excinfo:
            self.cls._get_limit_check_id()
        assert self.mock_conn.mock_calls == [
            call.describe_trusted_advisor_checks('en')
        ]
        assert excinfo.value.status == 400
        assert excinfo.value.reason == 'foobar'
        assert excinfo.value.body['__type'] == 'OtherException'

    def test_poll_id_none(self):
        tmp = self.mock_conn.describe_trusted_advisor_check_result
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            mock_id.return_value = None
            self.cls._poll()
        assert tmp.mock_calls == []

    def test_poll(self):
        tmp = self.mock_conn.describe_trusted_advisor_check_result
        tmp.return_value = {
            'result': {
                'timestamp': '2015-06-15T20:27:42Z',
                'flaggedResources': [
                    {
                        'status': 'ok',
                        'resourceId': 'resid1',
                        'isSuppressed': False,
                        'region': 'us-west-2',
                        'metadata': [
                            'us-west-2',
                            'AutoScaling',
                            'Auto Scaling groups',
                            '20',
                            '2',
                            'Green'
                        ]
                    },
                    {
                        'status': 'ok',
                        'resourceId': 'resid2',
                        'isSuppressed': False,
                        'region': 'us-east-1',
                        'metadata': [
                            'us-east-1',
                            'AutoScaling',
                            'Launch configurations',
                            '20',
                            '18',
                            'Yellow'
                        ]
                    },
                    {
                        'status': 'ok',
                        'resourceId': 'resid3',
                        'isSuppressed': False,
                        'region': 'us-east-1',
                        'metadata': [
                            'us-west-2',
                            'AutoScaling',
                            'Auto Scaling groups',
                            '40',
                            '10',
                            'Green'
                        ]
                    },
                ]
            }
        }
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            mock_id.return_value = (
                'foo',
                [
                    'Region',
                    'Service',
                    'Limit Name',
                    'Limit Amount',
                    'Current Usage',
                    'Status'
                ]
            )
            res = self.cls._poll()
        assert tmp.mock_calls == [call('foo')]
        assert mock_id.mock_calls == [call(self.cls)]
        assert res == {
            'AutoScaling': {
                'Launch configurations': 20,
                'Auto Scaling groups': 40,
            }
        }

    def test_poll_dont_have_ta(self):
        self.cls.have_ta = False
        tmp = self.mock_conn.describe_trusted_advisor_check_result
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            with patch('awslimitchecker.trustedadvisor'
                       '.logger', autospec=True) as mock_logger:
                res = self.cls._poll()
        assert tmp.mock_calls == []
        assert mock_id.mock_calls == [
            call(self.cls)
        ]
        assert mock_logger.mock_calls == [
            call.info('Beginning TrustedAdvisor poll'),
            call.info('TrustedAdvisor.have_ta is False; not polling TA')
        ]
        assert res == {}

    def test_update_services(self):

        def se_set(lname, val):
            if lname == 'blam':
                raise ValueError("foo")

        mock_autoscale = Mock(spec_set=_AwsService)
        mock_ec2 = Mock(spec_set=_AwsService)
        mock_ec2._set_ta_limit.side_effect = se_set
        services = {
            'AutoScaling': mock_autoscale,
            'EC2': mock_ec2,
        }
        ta_results = {
            'AutoScaling': {
                'foo': 20,
                'bar': 40,
            },
            'EC2': {
                'baz': 5,
                'blam': 10,
            },
            'OtherService': {
                'blarg': 1,
            }
        }
        with patch('awslimitchecker.trustedadvisor'
                   '.logger', autospec=True) as mock_logger:
            self.cls._update_services(ta_results, services)
        assert mock_logger.mock_calls == [
            call.debug("Updating TA limits on all services"),
            call.info("TrustedAdvisor returned check results for unknown "
                      "limit '%s' (service %s)", 'blam', 'EC2'),
            call.info("TrustedAdvisor returned check results for unknown "
                      "service '%s'", 'OtherService'),
            call.info("Done updating TA limits on all services"),
        ]
        assert mock_autoscale.mock_calls == [
            call._set_ta_limit('bar', 40),
            call._set_ta_limit('foo', 20),
        ]
        assert mock_ec2.mock_calls == [
            call._set_ta_limit('baz', 5),
            call._set_ta_limit('blam', 10),
        ]
class Test_TrustedAdvisor(object):
    def setup(self):
        self.mock_conn = Mock()
        self.mock_client_config = Mock()
        type(self.mock_client_config).region_name = 'us-east-1'
        type(self.mock_conn)._client_config = self.mock_client_config
        self.cls = TrustedAdvisor()
        self.cls.conn = self.mock_conn

        self.mock_svc1 = Mock(spec_set=_AwsService)
        self.mock_svc2 = Mock(spec_set=_AwsService)
        self.services = {
            'SvcFoo': self.mock_svc1,
            'SvcBar': self.mock_svc2,
        }

    def test_init(self):
        cls = TrustedAdvisor()
        assert cls.conn is None
        assert cls.account_id is None
        assert cls.account_role is None
        assert cls.region == 'us-east-1'
        assert cls.ta_region is None
        assert cls.external_id is None
        assert cls.mfa_serial_number is None
        assert cls.mfa_token is None

    def test_init_sts(self):
        cls = TrustedAdvisor(account_id='aid', account_role='role', region='r')
        assert cls.conn is None
        assert cls.account_id == 'aid'
        assert cls.account_role == 'role'
        assert cls.region == 'us-east-1'
        assert cls.ta_region == 'r'
        assert cls.external_id is None
        assert cls.mfa_serial_number is None
        assert cls.mfa_token is None

    def test_init_sts_external_id(self):
        cls = TrustedAdvisor(account_id='aid',
                             account_role='role',
                             region='r',
                             external_id='myeid')
        assert cls.conn is None
        assert cls.account_id == 'aid'
        assert cls.account_role == 'role'
        assert cls.region == 'us-east-1'
        assert cls.ta_region == 'r'
        assert cls.external_id == 'myeid'
        assert cls.mfa_serial_number is None
        assert cls.mfa_token is None

    def test_update_limits(self):
        mock_results = Mock()
        with patch('%s.connect' % pb, autospec=True) as mock_connect:
            with patch('%s._poll' % pb, autospec=True) as mock_poll:
                with patch('%s._update_services' % pb,
                           autospec=True) as mock_update_services:
                    mock_poll.return_value = mock_results
                    self.cls.update_limits(self.services)
        assert mock_connect.mock_calls == [call(self.cls)]
        assert mock_poll.mock_calls == [call(self.cls)]
        assert mock_update_services.mock_calls == [
            call(self.cls, mock_results, self.services)
        ]

    def test_get_limit_check_id(self):
        api_resp = {
            'checks': [
                {
                    'category':
                    'performance',
                    'name':
                    'Service Limits',
                    'id':
                    'bar',
                    'metadata': [
                        'Region', 'Service', 'Limit Name', 'Limit Amount',
                        'Current Usage', 'Status'
                    ],
                },
                {
                    'category': 'fault_tolerance',
                    'name': 'ELB Cross-Zone Load Balancing',
                    'id': 'foo',
                },
            ]
        }
        self.mock_conn.describe_trusted_advisor_checks.return_value = api_resp
        res = self.cls._get_limit_check_id()
        assert res == ('bar', [
            'Region', 'Service', 'Limit Name', 'Limit Amount', 'Current Usage',
            'Status'
        ])
        assert self.mock_conn.mock_calls == [
            call.describe_trusted_advisor_checks(language='en')
        ]

    def test_get_limit_check_id_none(self):
        api_resp = {
            'checks': [
                {
                    'category': 'performance',
                    'name': 'Something Else',
                    'id': 'bar',
                },
                {
                    'category': 'fault_tolerance',
                    'name': 'ELB Cross-Zone Load Balancing',
                    'id': 'foo',
                },
            ]
        }
        self.mock_conn.describe_trusted_advisor_checks.return_value = api_resp
        res = self.cls._get_limit_check_id()
        assert res == (None, None)
        assert self.mock_conn.mock_calls == [
            call.describe_trusted_advisor_checks(language='en')
        ]

    def test_get_limit_check_id_subscription_required(self):
        def se_api(language=None):
            response = {
                'ResponseMetadata': {
                    'HTTPStatusCode': 400,
                    'RequestId': '3cc9b2a8-c6e5-11e5-bc1d-b13dcea36176'
                },
                'Error': {
                    'Message': 'AWS Premium Support Subscription is required '
                    'to use this service.',
                    'Code': 'SubscriptionRequiredException'
                }
            }
            raise ClientError(response, 'operation')

        assert self.cls.have_ta is True
        self.mock_conn.describe_trusted_advisor_checks.side_effect = se_api
        with patch('awslimitchecker.trustedadvisor'
                   '.logger', autospec=True) as mock_logger:
            res = self.cls._get_limit_check_id()
        assert self.cls.have_ta is False
        assert res == (None, None)
        assert self.mock_conn.mock_calls == [
            call.describe_trusted_advisor_checks(language='en')
        ]
        assert mock_logger.mock_calls == [
            call.debug("Querying Trusted Advisor checks"),
            call.warning(
                "Cannot check TrustedAdvisor: %s",
                'AWS Premium Support Subscription is required to '
                'use this service.')
        ]

    def test_get_limit_check_id_other_exception(self):
        def se_api(language=None):
            response = {
                'ResponseMetadata': {
                    'HTTPStatusCode': 400,
                    'RequestId': '3cc9b2a8-c6e5-11e5-bc1d-b13dcea36176'
                },
                'Error': {
                    'Message': 'foo',
                    'Code': 'SomeOtherException'
                }
            }
            raise ClientError(response, 'operation')

        self.mock_conn.describe_trusted_advisor_checks.side_effect = se_api
        with pytest.raises(ClientError) as excinfo:
            self.cls._get_limit_check_id()
        assert self.mock_conn.mock_calls == [
            call.describe_trusted_advisor_checks(language='en')
        ]
        assert excinfo.value.response['ResponseMetadata'][
            'HTTPStatusCode'] == 400
        assert excinfo.value.response['Error']['Message'] == 'foo'
        assert excinfo.value.response['Error']['Code'] == 'SomeOtherException'

    def test_poll_id_none(self):
        tmp = self.mock_conn.describe_trusted_advisor_check_result
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            mock_id.return_value = None
            res = self.cls._poll()
        assert tmp.mock_calls == []
        assert res is None

    def test_poll(self):
        tmp = self.mock_conn.describe_trusted_advisor_check_result
        poll_return_val = {
            'result': {
                'timestamp':
                '2015-06-15T20:27:42Z',
                'flaggedResources': [
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid1',
                        'isSuppressed':
                        False,
                        'region':
                        'us-west-2',
                        'metadata': [
                            'us-west-2', 'AutoScaling', 'Auto Scaling groups',
                            '20', '2', 'Green'
                        ]
                    },
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid2',
                        'isSuppressed':
                        False,
                        'region':
                        'us-east-1',
                        'metadata': [
                            'us-east-1', 'AutoScaling',
                            'Launch configurations', '20', '18', 'Yellow'
                        ]
                    },
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid3',
                        'isSuppressed':
                        False,
                        'region':
                        'us-east-1',
                        'metadata': [
                            'us-west-2', 'AutoScaling', 'Auto Scaling groups',
                            '40', '10', 'Green'
                        ]
                    },
                ]
            }
        }
        tmp.return_value = poll_return_val
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            mock_id.return_value = ('foo', [
                'Region', 'Service', 'Limit Name', 'Limit Amount',
                'Current Usage', 'Status'
            ])
            res = self.cls._poll()
        assert self.mock_conn.mock_calls == [
            call.describe_trusted_advisor_check_result(checkId='foo',
                                                       language='en')
        ]
        assert mock_id.mock_calls == [call(self.cls)]
        assert res == {
            'AutoScaling': {
                'Launch configurations': 20,
                'Auto Scaling groups': 40,
            }
        }

    def test_poll_region(self):
        tmp = self.mock_conn.describe_trusted_advisor_check_result
        self.cls.ta_region = 'us-west-2'
        poll_return_value = {
            'result': {
                'timestamp':
                '2015-06-15T20:27:42Z',
                'flaggedResources': [
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid1',
                        'isSuppressed':
                        False,
                        'region':
                        'us-west-2',
                        'metadata': [
                            'us-west-2', 'AutoScaling', 'Auto Scaling groups',
                            '20', '2', 'Green'
                        ]
                    },
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid2',
                        'isSuppressed':
                        False,
                        'region':
                        'us-east-1',
                        'metadata': [
                            'us-east-1', 'AutoScaling',
                            'Launch configurations', '20', '18', 'Yellow'
                        ]
                    },
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid3',
                        'isSuppressed':
                        False,
                        'region':
                        'us-east-1',
                        'metadata': [
                            'us-west-2', 'AutoScaling', 'Auto Scaling groups',
                            '40', '10', 'Green'
                        ]
                    },
                ]
            }
        }
        tmp.return_value = poll_return_value
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            mock_id.return_value = ('foo', [
                'Region', 'Service', 'Limit Name', 'Limit Amount',
                'Current Usage', 'Status'
            ])
            res = self.cls._poll()
        assert self.mock_conn.mock_calls == [
            call.describe_trusted_advisor_check_result(checkId='foo',
                                                       language='en')
        ]
        assert mock_id.mock_calls == [call(self.cls)]
        assert res == {
            'AutoScaling': {
                'Auto Scaling groups': 20,
            }
        }

    def test_poll_dont_have_ta(self):
        self.cls.have_ta = False
        tmp = self.mock_conn.describe_trusted_advisor_check_result
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            with patch('awslimitchecker.trustedadvisor'
                       '.logger',
                       autospec=True) as mock_logger:
                res = self.cls._poll()
        assert self.mock_conn.mock_calls == []
        assert tmp.mock_calls == []
        assert mock_id.mock_calls == [call(self.cls)]
        assert mock_logger.mock_calls == [
            call.info('Beginning TrustedAdvisor poll'),
            call.info('TrustedAdvisor.have_ta is False; not polling TA')
        ]
        assert res == {}

    def test_update_services(self):
        def se_set(lname, val):
            if lname == 'blam':
                raise ValueError("foo")

        mock_autoscale = Mock(spec_set=_AwsService)
        mock_ec2 = Mock(spec_set=_AwsService)
        mock_ec2._set_ta_limit.side_effect = se_set
        mock_vpc = Mock(spec_set=_AwsService)
        services = {
            'AutoScaling': mock_autoscale,
            'EC2': mock_ec2,
            'VPC': mock_vpc,
        }
        ta_results = {
            'AutoScaling': {
                'foo': 20,
                'bar': 40,
            },
            'EC2': {
                'baz': 5,
                'blam': 10,
            },
            'OtherService': {
                'blarg': 1,
            },
            'VPC': {
                'VPC Elastic IP addresses (EIPs)': 11,
            }
        }
        with patch('awslimitchecker.trustedadvisor'
                   '.logger', autospec=True) as mock_logger:
            self.cls._update_services(ta_results, services)
        assert mock_logger.mock_calls == [
            call.debug("Updating TA limits on all services"),
            call.info(
                "TrustedAdvisor returned check results for unknown "
                "limit '%s' (service %s)", 'blam', 'EC2'),
            call.info(
                "TrustedAdvisor returned check results for unknown "
                "service '%s'", 'OtherService'),
            call.info("Done updating TA limits on all services"),
        ]
        assert mock_autoscale.mock_calls == [
            call._set_ta_limit('bar', 40),
            call._set_ta_limit('foo', 20),
        ]
        assert mock_ec2.mock_calls == [
            call._set_ta_limit('baz', 5),
            call._set_ta_limit('blam', 10),
            call._set_ta_limit('VPC Elastic IP addresses (EIPs)', 11)
        ]

    def test_update_services_no_ec2(self):

        mock_autoscale = Mock(spec_set=_AwsService)
        mock_vpc = Mock(spec_set=_AwsService)
        services = {
            'AutoScaling': mock_autoscale,
            'VPC': mock_vpc,
        }
        ta_results = {
            'AutoScaling': {
                'foo': 20,
                'bar': 40,
            },
            'EC2': {
                'baz': 5,
            },
            'VPC': {
                'VPC Elastic IP addresses (EIPs)': 11,
            }
        }
        with patch('awslimitchecker.trustedadvisor'
                   '.logger', autospec=True) as mock_logger:
            self.cls._update_services(ta_results, services)
        assert mock_logger.mock_calls == [
            call.debug("Updating TA limits on all services"),
            call.info(
                "TrustedAdvisor returned check results for unknown "
                "service '%s'", 'EC2'),
            call.info("Done updating TA limits on all services"),
        ]
        assert mock_autoscale.mock_calls == [
            call._set_ta_limit('bar', 40),
            call._set_ta_limit('foo', 20),
        ]
        assert mock_vpc.mock_calls == []
Esempio n. 5
0
class TestPoll(object):
    def setup(self):
        self.mock_conn = Mock()
        self.mock_client_config = Mock()
        type(self.mock_client_config).region_name = 'us-east-1'
        type(self.mock_conn)._client_config = self.mock_client_config
        self.cls = TrustedAdvisor({}, {})
        self.cls.conn = self.mock_conn

        self.mock_svc1 = Mock(spec_set=_AwsService)
        self.mock_svc2 = Mock(spec_set=_AwsService)
        self.services = {
            'SvcFoo': self.mock_svc1,
            'SvcBar': self.mock_svc2,
        }

    def test_none(self):
        tmp = self.mock_conn.describe_trusted_advisor_check_result
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            mock_id.return_value = None
            res = self.cls._poll()
        assert tmp.mock_calls == []
        assert res is None

    def test_basic(self):
        poll_return_val = {
            'result': {
                'timestamp':
                '2015-06-15T20:27:42Z',
                'flaggedResources': [
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid1',
                        'isSuppressed':
                        False,
                        'region':
                        'us-west-2',
                        'metadata': [
                            'us-west-2', 'AutoScaling', 'Auto Scaling groups',
                            '20', '2', 'Green'
                        ]
                    },
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid1',
                        'isSuppressed':
                        False,
                        'region':
                        'us-east-1',
                        'metadata': [
                            'us-east-1', 'EC2',
                            'On-Demand instances - t2.micro', 'Unlimited', '2',
                            'Green'
                        ]
                    },
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid1',
                        'isSuppressed':
                        False,
                        'region':
                        'us-east-1',
                        'metadata': [
                            'us-east-1', 'EC2',
                            'On-Demand instances - t2.small', 'error', '2',
                            'Green'
                        ]
                    },
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid2',
                        'isSuppressed':
                        False,
                        'region':
                        'us-east-1',
                        'metadata': [
                            'us-east-1', 'AutoScaling',
                            'Launch configurations', '20', '18', 'Yellow'
                        ]
                    },
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid3',
                        'isSuppressed':
                        False,
                        'region':
                        'us-east-1',
                        'metadata': [
                            'us-west-2', 'AutoScaling', 'Auto Scaling groups',
                            '40', '10', 'Green'
                        ]
                    },
                    {
                        'status': 'ok',
                        'resourceId': 'resid4',
                        'isSuppressed': False,
                        'metadata':
                        ['-', 'IAM', 'Users', '5000', '152', 'Green']
                    },
                ]
            }
        }
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            with patch('%s._get_refreshed_check_result' % pb,
                       autospec=True) as mock_hr:
                mock_hr.return_value = poll_return_val
                mock_id.return_value = ('foo', [
                    'Region', 'Service', 'Limit Name', 'Limit Amount',
                    'Current Usage', 'Status'
                ])
                res = self.cls._poll()
        assert self.mock_conn.mock_calls == []
        assert mock_id.mock_calls == [call(self.cls)]
        assert mock_hr.mock_calls == [call(self.cls, 'foo')]
        assert res == {
            'AutoScaling': {
                'Launch configurations': 20,
                'Auto Scaling groups': 40,
            },
            'EC2': {
                'On-Demand instances - t2.micro': 'Unlimited'
            },
            'IAM': {
                'Users': 5000
            }
        }

    def test_no_timestamp(self):
        poll_return_val = {
            'result': {
                'flaggedResources': [
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid1',
                        'isSuppressed':
                        False,
                        'region':
                        'us-west-2',
                        'metadata': [
                            'us-west-2', 'AutoScaling', 'Auto Scaling groups',
                            '20', '2', 'Green'
                        ]
                    },
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid1',
                        'isSuppressed':
                        False,
                        'region':
                        'us-east-1',
                        'metadata': [
                            'us-east-1', 'EC2',
                            'On-Demand instances - t2.micro', 'Unlimited', '2',
                            'Green'
                        ]
                    },
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid1',
                        'isSuppressed':
                        False,
                        'region':
                        'us-east-1',
                        'metadata': [
                            'us-east-1', 'EC2',
                            'On-Demand instances - t2.small', 'error', '2',
                            'Green'
                        ]
                    },
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid2',
                        'isSuppressed':
                        False,
                        'region':
                        'us-east-1',
                        'metadata': [
                            'us-east-1', 'AutoScaling',
                            'Launch configurations', '20', '18', 'Yellow'
                        ]
                    },
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid3',
                        'isSuppressed':
                        False,
                        'region':
                        'us-east-1',
                        'metadata': [
                            'us-west-2', 'AutoScaling', 'Auto Scaling groups',
                            '40', '10', 'Green'
                        ]
                    },
                    {
                        'status': 'ok',
                        'resourceId': 'resid4',
                        'isSuppressed': False,
                        'metadata':
                        ['-', 'IAM', 'Users', '5000', '152', 'Green']
                    },
                ]
            }
        }
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            with patch('%s._get_refreshed_check_result' % pb,
                       autospec=True) as mock_hr:
                mock_hr.return_value = poll_return_val
                mock_id.return_value = ('foo', [
                    'Region', 'Service', 'Limit Name', 'Limit Amount',
                    'Current Usage', 'Status'
                ])
                res = self.cls._poll()
        assert self.mock_conn.mock_calls == []
        assert mock_id.mock_calls == [call(self.cls)]
        assert mock_hr.mock_calls == [call(self.cls, 'foo')]
        assert res == {
            'AutoScaling': {
                'Launch configurations': 20,
                'Auto Scaling groups': 40,
            },
            'EC2': {
                'On-Demand instances - t2.micro': 'Unlimited'
            },
            'IAM': {
                'Users': 5000
            }
        }

    def test_region(self):
        self.cls.ta_region = 'us-west-2'
        poll_return_value = {
            'result': {
                'timestamp':
                '2015-06-15T20:27:42Z',
                'flaggedResources': [
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid1',
                        'isSuppressed':
                        False,
                        'region':
                        'us-west-2',
                        'metadata': [
                            'us-west-2', 'AutoScaling', 'Auto Scaling groups',
                            '20', '2', 'Green'
                        ]
                    },
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid2',
                        'isSuppressed':
                        False,
                        'region':
                        'us-east-1',
                        'metadata': [
                            'us-east-1', 'AutoScaling',
                            'Launch configurations', '20', '18', 'Yellow'
                        ]
                    },
                    {
                        'status':
                        'ok',
                        'resourceId':
                        'resid3',
                        'isSuppressed':
                        False,
                        'region':
                        'us-east-1',
                        'metadata': [
                            'us-west-2', 'AutoScaling', 'Auto Scaling groups',
                            '40', '10', 'Green'
                        ]
                    },
                    {
                        'status': 'ok',
                        'resourceId': 'resid4',
                        'isSuppressed': False,
                        'metadata':
                        ['-', 'IAM', 'Users', '5000', '152', 'Green']
                    },
                ]
            }
        }
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            with patch('%s._get_refreshed_check_result' % pb,
                       autospec=True) as mock_hr:
                mock_hr.return_value = poll_return_value
                mock_id.return_value = ('foo', [
                    'Region', 'Service', 'Limit Name', 'Limit Amount',
                    'Current Usage', 'Status'
                ])
                res = self.cls._poll()
        assert self.mock_conn.mock_calls == []
        assert mock_id.mock_calls == [call(self.cls)]
        assert mock_hr.mock_calls == [call(self.cls, 'foo')]
        assert res == {
            'AutoScaling': {
                'Auto Scaling groups': 20,
            },
            'IAM': {
                'Users': 5000
            }
        }

    def test_dont_have_ta(self):
        self.cls.have_ta = False
        with patch('%s._get_limit_check_id' % pb, autospec=True) as mock_id:
            with patch('%s._get_refreshed_check_result' % pb,
                       autospec=True) as mock_hr:
                with patch('awslimitchecker.trustedadvisor'
                           '.logger',
                           autospec=True) as mock_logger:
                    res = self.cls._poll()
        assert self.mock_conn.mock_calls == []
        assert mock_id.mock_calls == [call(self.cls)]
        assert mock_hr.mock_calls == []
        assert mock_logger.mock_calls == [
            call.info('Beginning TrustedAdvisor poll'),
            call.info('TrustedAdvisor.have_ta is False; not polling TA')
        ]
        assert res == {}