PR #64 - add unit tests for STS in checker and runner

awslimitchecker/checker.py

@@ -66,9 +66,13 @@ def __init__(self, warning_threshold=80, critical_threshold=99,
           integer percentage, for any limits without a specifically-set
           threshold.
         :type critical_threshold: int
-        :param account_id: connect via STS to this AWS account
+        :param account_id: `AWS Account ID <http://docs.aws.amazon.com/general/latest/gr/acct-identifiers.html>`_
+          (12-digit string, currently numeric) for the account to connect to
+          (destination) via STS
         :type account_id: str
-        :param account_role: connect via STS as this IAM role
+        :param account_role: the name of an
+          `IAM Role <http://docs.aws.amazon.com/IAM/latest/UserGuide/id_roles.html>`_
+          (in the destination account) to assume
         :type account_role: str
         """
         # ###### IMPORTANT license notice ##########
@@ -100,7 +104,11 @@ def __init__(self, warning_threshold=80, critical_threshold=99,
         self.account_role = account_role
         self.region = region
         self.services = {}
-        self.ta = TrustedAdvisor(account_id=self.account_id, account_role=self.account_role, region=self.region)
+        self.ta = TrustedAdvisor(
+            account_id=self.account_id,
+            account_role=self.account_role,
+            region=self.region
+        )
         for sname, cls in _services.items():
             self.services[sname] = cls(warning_threshold, critical_threshold,
                                        account_id, account_role, region)

awslimitchecker/runner.py

@@ -130,13 +130,15 @@ def parse_args(self, argv):
                        'limit); default: 99')
         p.add_argument('-A', '--sts-account-id', action='store',
                        type=str, default=None,
-                       help='the AWS account to control')
+                       help='for use with STS, the Account ID of the '
+                       'destination account (account to assume a role in)')
         p.add_argument('-R', '--sts-account-role', action='store',
                        type=str, default=None,
-                       help='the IAM role to assume')
+                       help='for use with STS, the name of the IAM role to '
+                       'assume')
         p.add_argument('-r', '--region', action='store',
                        type=str, default=None,
-                       help='connect to this AWS region; required for STS')
+                       help='AWS region name to connect to; required for STS')
         p.add_argument('--skip-ta', action='store_true', default=False,
                        help='do not attempt to pull *any* information on limits'
                        ' from Trusted Advisor')

awslimitchecker/tests/test_checker.py

@@ -88,6 +88,7 @@ def setup(self):
             ) 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()
@@ -101,6 +102,9 @@ def test_init(self):
         # _AwsService instances should exist, but have no other calls
         assert self.mock_foo.mock_calls == [call(80, 99, None, None, None)]
         assert self.mock_bar.mock_calls == [call(80, 99, None, None, None)]
+        assert self.mock_ta_constr.mock_calls == [
+            call(account_id=None, account_role=None, region=None)
+        ]
         assert self.mock_svc1.mock_calls == []
         assert self.mock_svc2.mock_calls == []
         assert self.cls.ta == self.mock_ta
@@ -128,6 +132,7 @@ def test_init_thresholds(self):
         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}
@@ -137,10 +142,13 @@ def test_init_thresholds(self):
                     '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 = self.mock_ta
                 cls = AwsLimitChecker(
                     warning_threshold=5,
                     critical_threshold=22,
@@ -153,6 +161,98 @@ def test_init_thresholds(self):
         # _AwsService instances should exist, but have no other calls
         assert mock_foo.mock_calls == [call(5, 22, None, None, None)]
         assert mock_bar.mock_calls == [call(5, 22, None, None, None)]
+        assert mock_ta_constr.mock_calls == [
+            call(account_id=None, account_role=None, region=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 = self.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')]
+        assert mock_bar.mock_calls == [call(80, 99, None, None, 'myregion')]
+        assert mock_ta_constr.mock_calls == [
+            call(account_id=None, account_role=None, region='myregion')
+        ]
+        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 = self.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')
+        ]
+        assert mock_bar.mock_calls == [
+            call(80, 99, '123456789012', 'myrole', 'myregion')
+        ]
+        assert mock_ta_constr.mock_calls == [
+            call(
+                account_id='123456789012',
+                account_role='myrole',
+                region='myregion')
+        ]
         assert mock_svc1.mock_calls == []
         assert mock_svc2.mock_calls == []
         assert self.mock_version.mock_calls == [call()]

awslimitchecker/tests/test_runner.py

@@ -159,14 +159,17 @@ def test_parse_args_parser(self):
                                 'of limit); default: 99'),
             call().add_argument('-A', '--sts-account-id', action='store',
                                 type=str, default=None,
-                                help='the AWS account to control'),
+                                help='for use with STS, the Account ID of the '
+                                'destination account (account to assume a role'
+                                ' in)'),
             call().add_argument('-R', '--sts-account-role', action='store',
                                 type=str, default=None,
-                                help='the IAM role to assume'),
+                                help='for use with STS, the name of the IAM '
+                                'role to assume'),
             call().add_argument('-r', '--region', action='store',
                                 type=str, default=None,
-                                help='connect to this AWS region; '
-                                'required for STS'),
+                                help='AWS region name to connect to; required '
+                                'for STS'),
             call().add_argument('--skip-ta', action='store_true', default=False,
                                 help='do not attempt to pull *any* information '
                                 'on limits from Trusted Advisor'),
@@ -532,6 +535,62 @@ def test_entry_no_service_name(self, capsys):
         assert excinfo.value.code == 6
         assert self.cls.service_name is None
 
+    def test_entry_no_service_name_region(self, capsys):
+        argv = ['awslimitchecker', '-r', 'myregion']
+        with patch.object(sys, 'argv', argv):
+            with patch('%s.Runner.check_thresholds' % pb,
+                       autospec=True) as mock_ct:
+                with patch('%s.AwsLimitChecker' % pb,
+                           spec_set=AwsLimitChecker) as mock_alc:
+                    with pytest.raises(SystemExit) as excinfo:
+                        mock_ct.return_value = 6
+                        self.cls.console_entry_point()
+        out, err = capsys.readouterr()
+        assert out == ''
+        assert excinfo.value.code == 6
+        assert mock_alc.mock_calls == [
+            call(
+                warning_threshold=80,
+                critical_threshold=99,
+                account_id=None,
+                account_role=None,
+                region='myregion'
+            )
+        ]
+        assert self.cls.service_name is None
+
+    def test_entry_no_service_name_sts(self, capsys):
+        argv = [
+            'awslimitchecker',
+            '-r',
+            'myregion',
+            '-A',
+            '098765432109',
+            '-R',
+            'myrole'
+        ]
+        with patch.object(sys, 'argv', argv):
+            with patch('%s.Runner.check_thresholds' % pb,
+                       autospec=True) as mock_ct:
+                with patch('%s.AwsLimitChecker' % pb,
+                           spec_set=AwsLimitChecker) as mock_alc:
+                    with pytest.raises(SystemExit) as excinfo:
+                        mock_ct.return_value = 6
+                        self.cls.console_entry_point()
+        out, err = capsys.readouterr()
+        assert out == ''
+        assert excinfo.value.code == 6
+        assert mock_alc.mock_calls == [
+            call(
+                warning_threshold=80,
+                critical_threshold=99,
+                account_id='098765432109',
+                account_role='myrole',
+                region='myregion'
+            )
+        ]
+        assert self.cls.service_name is None
+
     def test_entry_verbose(self, capsys):
         argv = ['awslimitchecker', '-v']
         with patch.object(sys, 'argv', argv):