Add methods for the persistent cache Agent interface (#6819)

DataDog · Jun 4, 2020 · 3f4d637 · 3f4d637
1 parent 3669349
commit 3f4d637
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 0 deletions.
diff --git 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
@@ -15,13 +15,15 @@ class DatadogAgentStub(object):
 
     def __init__(self):
         self._metadata = {}
+        self._cache = {}
         self._config = self.get_default_config()
 
     def get_default_config(self):
         return {'enable_metadata_collection': True}
 
     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
@@ -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)],