From b9364fc08115570145d2ae4410f1870ad50fa66c Mon Sep 17 00:00:00 2001 From: Ofek Lev Date: Thu, 4 Jun 2020 02:12:09 -0400 Subject: [PATCH] Add methods for the persistent cache Agent interface --- .../datadog_checks/base/checks/base.py | 12 ++++++++++++ .../datadog_checks/base/stubs/datadog_agent.py | 8 ++++++++ datadog_checks_base/tests/test_agent_check.py | 10 ++++++++++ 3 files changed, 30 insertions(+) diff --git a/datadog_checks_base/datadog_checks/base/checks/base.py b/datadog_checks_base/datadog_checks/base/checks/base.py index c44daed603495..ca67f19e2b1b5 100644 --- a/datadog_checks_base/datadog_checks/base/checks/base.py +++ b/datadog_checks_base/datadog_checks/base/checks/base.py @@ -657,6 +657,18 @@ def entrypoint(self, *args, **kwargs): return entrypoint + def _persistent_cache_id(self, key): + # type: (str) -> str + return '{}_{}'.format(self.check_id, key) + + def read_persistent_cache(self, key): + # type: (str) -> str + return datadog_agent.read_persistent_cache(self._persistent_cache_id(key)) + + def write_persistent_cache(self, key, value): + # type: (str, str) -> None + datadog_agent.write_persistent_cache(self._persistent_cache_id(key), value) + def set_external_tags(self, external_tags): # type: (Sequence[ExternalTagType]) -> None # Example of external_tags format diff --git a/datadog_checks_base/datadog_checks/base/stubs/datadog_agent.py b/datadog_checks_base/datadog_checks/base/stubs/datadog_agent.py index c49d13cac4767..6e9620bc55447 100644 --- a/datadog_checks_base/datadog_checks/base/stubs/datadog_agent.py +++ b/datadog_checks_base/datadog_checks/base/stubs/datadog_agent.py @@ -15,6 +15,7 @@ class DatadogAgentStub(object): def __init__(self): self._metadata = {} + self._cache = {} self._config = self.get_default_config() def get_default_config(self): @@ -22,6 +23,7 @@ def get_default_config(self): def reset(self): self._metadata.clear() + self._cache.clear() self._config = self.get_default_config() def assert_metadata(self, check_id, data): @@ -56,6 +58,12 @@ def set_external_tags(self, *args, **kwargs): def tracemalloc_enabled(self, *args, **kwargs): return False + def write_persistent_cache(self, key, value): + self._cache[key] = value + + def read_persistent_cache(self, key): + return self._cache.get(key, '') + # Use the stub as a singleton datadog_agent = DatadogAgentStub() diff --git a/datadog_checks_base/tests/test_agent_check.py b/datadog_checks_base/tests/test_agent_check.py index f0515c101f3da..964761afe6fa4 100644 --- a/datadog_checks_base/tests/test_agent_check.py +++ b/datadog_checks_base/tests/test_agent_check.py @@ -44,6 +44,16 @@ def test_load_config(): assert AgentCheck.load_config("raw_foo: bar") == {'raw_foo': 'bar'} +def test_persistent_cache(datadog_agent): + check = AgentCheck() + check.check_id = 'test' + + check.write_persistent_cache('foo', 'bar') + + assert datadog_agent.read_persistent_cache('test_foo') == 'bar' + assert check.read_persistent_cache('foo') == 'bar' + + @pytest.mark.parametrize( 'enable_metadata_collection, expected_is_metadata_collection_enabled', [(None, False), ('true', True), ('false', False)],