fixes #11 - implement threshold overrides and expose up to Checker

awslimitchecker/checker.py

@@ -223,6 +223,92 @@ def set_limit_override(self, service_name, limit_name,
             override_ta=override_ta
         )
 
+    def set_threshold_overrides(self, override_dict):
+        """
+        Set manual overrides on the threshold (used for determining
+        warning/critical status) a dict of limits. See
+        :py:class:`~.AwsLimitChecker` for information on Warning and
+        Critical thresholds.
+
+        Dict is composed of service name keys (string) to dict of
+        limit names (string), to dict of threshold specifications.
+        Each threhold specification dict can contain keys 'warning'
+        or 'critical', each having a value of a dict containing
+        keys 'percent' or 'count', to an integer value.
+
+        Example:
+
+            {
+                'EC2': {
+                    'SomeLimit': {
+                        'warning': {
+                            'percent': 80,
+                            'count': 8,
+                        },
+                        'critical': {
+                            'percent': 90,
+                            'count': 9,
+                        }
+                    }
+                }
+            }
+
+        See :py:meth:`~.AwsLimit.set_threshold_override`.
+
+        :param override_dict: nested dict of threshold overrides
+        :type override_dict: dict
+        """
+        for svc_name in sorted(override_dict):
+            for lim_name in sorted(override_dict[svc_name]):
+                d = override_dict[svc_name][lim_name]
+                kwargs = {}
+                if 'warning' in d:
+                    if 'percent' in d['warning']:
+                        kwargs['warn_percent'] = d['warning']['percent']
+                    if 'count' in d['warning']:
+                        kwargs['warn_count'] = d['warning']['count']
+                if 'critical' in d:
+                    if 'percent' in d['critical']:
+                        kwargs['crit_percent'] = d['critical']['percent']
+                    if 'count' in d['critical']:
+                        kwargs['crit_count'] = d['critical']['count']
+                self.services[svc_name].set_threshold_override(
+                    lim_name,
+                    **kwargs
+                )
+
+    def set_threshold_override(self, service_name, limit_name,
+                               warn_percent=None, warn_count=None,
+                               crit_percent=None, crit_count=None):
+        """
+        Set a manual override on the threshold (used for determining
+        warning/critical status) for a specific limit. See
+        :py:class:`~.AwsLimitChecker` for information on Warning and
+        Critical thresholds.
+
+        See :py:meth:`~.AwsLimit.set_threshold_override`.
+
+        :param service_name: the name of the service to override limit for
+        :type service_name: string
+        :param limit_name: the name of the limit to override:
+        :type limit_name: string
+        :param warn_percent: new warning threshold, percentage used
+        :type warn_percent: int
+        :param warn_count: new warning threshold, actual count/number
+        :type warn_count: int
+        :param crit_percent: new critical threshold, percentage used
+        :type crit_percent: int
+        :param crit_count: new critical threshold, actual count/number
+        :type crit_count: int
+        """
+        self.services[service_name].set_threshold_override(
+            limit_name,
+            warn_percent=warn_percent,
+            warn_count=warn_count,
+            crit_percent=crit_percent,
+            crit_count=crit_count
+        )
+
     def check_thresholds(self, service=None):
         """
         Check all limits and current usage against their specified thresholds;

awslimitchecker/limit.py

@@ -83,6 +83,10 @@ def __init__(self, name, service, default_limit,
         self._current_usage = []
         self.def_warning_threshold = def_warning_threshold
         self.def_critical_threshold = def_critical_threshold
+        self.warn_percent = None
+        self.warn_count = None
+        self.crit_percent = None
+        self.crit_count = None
         self._warnings = []
         self._criticals = []
 
@@ -201,15 +205,35 @@ def _get_thresholds(self):
 
         :rtype: tuple
         """
-        # @TODO threshold overrides
         t = (
-            None,
-            self.def_warning_threshold,
-            None,
-            self.def_critical_threshold,
+            self.warn_count,
+            self.warn_percent or self.def_warning_threshold,
+            self.crit_count,
+            self.crit_percent or self.def_critical_threshold,
         )
         return t
 
+    def set_threshold_override(self, warn_percent=None, warn_count=None,
+                               crit_percent=None, crit_count=None):
+        """
+        Override the default warning and critical thresholds used to evaluate
+        this limit's usage. Theresholds can be specified as a percentage
+        of the limit, or as a usage count, or both.
+
+        :param warn_percent: new warning threshold, percentage used
+        :type warn_percent: int
+        :param warn_count: new warning threshold, actual count/number
+        :type warn_count: int
+        :param crit_percent: new critical threshold, percentage used
+        :type crit_percent: int
+        :param crit_count: new critical threshold, actual count/number
+        :type crit_count: int
+        """
+        self.warn_percent = warn_percent
+        self.warn_count = warn_count
+        self.crit_percent = crit_percent
+        self.crit_count = crit_count
+
     def check_thresholds(self):
         """
         Check this limit's current usage against the specified default

awslimitchecker/services/base.py

@@ -174,6 +174,35 @@ def set_limit_override(self, limit_name, value, override_ta=True):
                 s=self.service_name,
                 l=limit_name))
 
+    def set_threshold_override(self, limit_name, warn_percent=None,
+                               warn_count=None, crit_percent=None,
+                               crit_count=None):
+        """
+        Override the default warning and critical thresholds used to evaluate
+        the specified limit's usage. Theresholds can be specified as a
+        percentage of the limit, or as a usage count, or both.
+
+        :param warn_percent: new warning threshold, percentage used
+        :type warn_percent: int
+        :param warn_count: new warning threshold, actual count/number
+        :type warn_count: int
+        :param crit_percent: new critical threshold, percentage used
+        :type crit_percent: int
+        :param crit_count: new critical threshold, actual count/number
+        :type crit_count: int
+        """
+        try:
+            self.limits[limit_name].set_threshold_override(
+                warn_percent=warn_percent,
+                warn_count=warn_count,
+                crit_percent=crit_percent,
+                crit_count=crit_count
+            )
+        except KeyError:
+            raise ValueError("{s} service has no '{l}' limit".format(
+                s=self.service_name,
+                l=limit_name))
+
     def check_thresholds(self):
         """
         Checks current usage against configured thresholds for all limits

awslimitchecker/tests/services/test_base.py

@@ -128,6 +128,44 @@ def test_set_limit_override_keyerror(self):
             "'bar' limit"
         assert mock_limit.mock_calls == []
 
+    def test_set_threshold_override(self):
+        mock_limit = Mock(spec_set=AwsLimit)
+        type(mock_limit).default_limit = 5
+        cls = AwsServiceTester(1, 2)
+        cls.limits['foo'] = mock_limit
+        cls.set_threshold_override(
+            'foo',
+            warn_percent=10,
+            warn_count=12,
+            crit_percent=14,
+            crit_count=16
+        )
+        assert mock_limit.mock_calls == [
+            call.set_threshold_override(
+                warn_percent=10,
+                warn_count=12,
+                crit_percent=14,
+                crit_count=16
+            )
+        ]
+
+    def test_set_threshold_override_keyerror(self):
+        mock_limit = Mock(spec_set=AwsLimit)
+        type(mock_limit).default_limit = 5
+        cls = AwsServiceTester(1, 2)
+        cls.limits['foo'] = mock_limit
+        with pytest.raises(ValueError) as excinfo:
+            cls.set_threshold_override('bar', warn_percent=10)
+
+        if sys.version_info[0] > 2:
+            msg = excinfo.value.args[0]
+        else:
+            msg = excinfo.value.message
+
+        assert msg == "AwsServiceTester service has no " \
+            "'bar' limit"
+        assert mock_limit.mock_calls == []
+
     def test_check_thresholds(self):
         cls = AwsServiceTester(1, 2)
         cls.find_usage()

awslimitchecker/tests/test_checker.py

@@ -41,6 +41,7 @@
 from awslimitchecker.services.base import _AwsService
 from awslimitchecker.checker import AwsLimitChecker
 from awslimitchecker.version import _get_version, _get_project_url
+from awslimitchecker.limit import AwsLimit
 from .support import sample_limits
 
 
@@ -176,6 +177,77 @@ def test_find_usage_service(self):
         ]
         assert self.mock_svc2.mock_calls == []
 
+    def test_set_threshold_overrides(self):
+        limits = sample_limits()
+        limits['SvcFoo']['zz3'] = AwsLimit(
+            'zz3',
+            self.mock_svc1,
+            1,
+            2,
+            3,
+        )
+        self.mock_svc1.get_limits.return_value = limits['SvcFoo']
+        self.mock_svc2.get_limits.return_value = limits['SvcBar']
+        overrides = {
+            'SvcBar': {
+                'barlimit1': {
+                    'warning': {
+                        'percent': 10,
+                        'count': 12
+                    },
+                    'critical': {
+                        'percent': 14,
+                        'count': 16
+                    }
+                },
+                'bar limit2': {
+                    'critical': {
+                        'count': 15,
+                    }
+                },
+                'zz3': {
+                    'warning': {
+                        'count': 41
+                    },
+                    'critical': {
+                        'percent': 52
+                    }
+                }
+            },
+            'SvcFoo': {
+                'foo limit3': {
+                    'warning': {
+                        'percent': 91
+                    },
+                }
+            },
+        }
+        self.cls.set_threshold_overrides(overrides)
+        assert self.mock_svc1.mock_calls == [
+            call.set_threshold_override(
+                'foo limit3',
+                warn_percent=91,
+            )
+        ]
+        assert self.mock_svc2.mock_calls == [
+            call.set_threshold_override(
+                'bar limit2',
+                crit_count=15
+            ),
+            call.set_threshold_override(
+                'barlimit1',
+                warn_percent=10,
+                warn_count=12,
+                crit_percent=14,
+                crit_count=16
+            ),
+            call.set_threshold_override(
+                'zz3',
+                warn_count=41,
+                crit_percent=52
+            ),
+        ]
+
     def test_set_limit_overrides(self):
         limits = sample_limits()
         self.mock_svc1.get_limits.return_value = limits['SvcFoo']
@@ -241,6 +313,27 @@ def test_set_limit_override_ta(self):
             )
         ]
 
+    def test_set_threshold_override(self):
+        limits = sample_limits()
+        self.mock_svc1.get_limits.return_value = limits['SvcFoo']
+        self.cls.set_threshold_override(
+            'SvcFoo',
+            'foo limit3',
+            warn_percent=10,
+            warn_count=12,
+            crit_percent=14,
+            crit_count=16
+        )
+        assert self.mock_svc1.mock_calls == [
+            call.set_threshold_override(
+                'foo limit3',
+                warn_percent=10,
+                warn_count=12,
+                crit_percent=14,
+                crit_count=16
+            )
+        ]
+
     def test_get_required_iam_policy(self):
         expected = {
             'Version': '2012-10-17',

awslimitchecker/tests/test_limit.py

@@ -378,6 +378,32 @@ def test_get_thresholds(self):
             2
         )
 
+    def test_get_thresholds_overridden(self):
+        limit = AwsLimit('limitname', self.mock_svc, 100, 88, 99)
+        limit.warn_percent = 1
+        limit.warn_count = 2
+        limit.crit_percent = 3
+        limit.crit_count = 4
+        assert limit._get_thresholds() == (
+            2,
+            1,
+            4,
+            3
+        )
+
+    def test_set_threshold_override(self):
+        limit = AwsLimit('limitname', self.mock_svc, 100, 1, 2)
+        limit.set_threshold_override(
+            warn_percent=1,
+            warn_count=2,
+            crit_percent=3,
+            crit_count=4
+        )
+        assert limit.warn_percent == 1
+        assert limit.warn_count == 2
+        assert limit.crit_percent == 3
+        assert limit.crit_count == 4
+
 
 class TestAwsLimitUsage(object):