diff --git a/datadog_checks_base/datadog_checks/base/checks/base.py b/datadog_checks_base/datadog_checks/base/checks/base.py index a55f3caca7805..11412df511a61 100644 --- a/datadog_checks_base/datadog_checks/base/checks/base.py +++ b/datadog_checks_base/datadog_checks/base/checks/base.py @@ -9,7 +9,7 @@ import re import traceback import unicodedata -from collections import defaultdict +from collections import defaultdict, deque from os.path import basename import yaml @@ -212,6 +212,9 @@ class except the :py:meth:`check` method but sometimes it might be useful for a if metric_limit > 0: self.metric_limiter = Limiter(self.name, 'metrics', metric_limit, self.warning) + # Functions that will be called exactly once (if successful) before the first check run + self.check_initializations = deque() + @staticmethod def load_config(yaml_str): """ @@ -600,6 +603,14 @@ def check(self, instance): def run(self): try: + while self.check_initializations: + initialization = self.check_initializations.popleft() + try: + initialization() + except Exception: + self.check_initializations.appendleft(initialization) + raise + instance = copy.deepcopy(self.instances[0]) if 'set_breakpoint' in self.init_config: diff --git a/datadog_checks_base/tests/test_agent_check.py b/datadog_checks_base/tests/test_agent_check.py index 089ddb9f198b8..483d97b3a2ab9 100644 --- a/datadog_checks_base/tests/test_agent_check.py +++ b/datadog_checks_base/tests/test_agent_check.py @@ -376,3 +376,51 @@ def test_metric_limit_instance_config_zero(self, aggregator): check.gauge("metric", 0) assert len(check.get_warnings()) == 1 # get_warnings resets the array assert len(aggregator.metrics("metric")) == 10 + + +class TestCheckInitializations: + def test_success_only_once(self): + class TestCheck(AgentCheck): + def __init__(self, *args, **kwargs): + super(TestCheck, self).__init__(*args, **kwargs) + self.state = 1 + self.initialize = mock.MagicMock(side_effect=self._initialize) + self.check_initializations.append(self.initialize) + + def _initialize(self): + self.state += 1 + if self.state % 2: + raise Exception('is odd') + + def check(self, _): + pass + + check = TestCheck('test', {}, [{}]) + check.run() + check.run() + check.run() + + assert check.initialize.call_count == 1 + + def test_error_retry(self): + class TestCheck(AgentCheck): + def __init__(self, *args, **kwargs): + super(TestCheck, self).__init__(*args, **kwargs) + self.state = 0 + self.initialize = mock.MagicMock(side_effect=self._initialize) + self.check_initializations.append(self.initialize) + + def _initialize(self): + self.state += 1 + if self.state % 2: + raise Exception('is odd') + + def check(self, _): + pass + + check = TestCheck('test', {}, [{}]) + check.run() + check.run() + check.run() + + assert check.initialize.call_count == 2