Add @metadata_entrypoint decorator

datadog_checks_base/datadog_checks/base/checks/base.py

@@ -2,6 +2,7 @@
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
 import copy
+import functools
 import importlib
 import inspect
 import json
@@ -634,6 +635,26 @@ def is_metadata_collection_enabled():
         # type: () -> bool
         return is_affirmative(datadog_agent.get_config('enable_metadata_collection'))
 
+    @classmethod
+    def metadata_entrypoint(cls, method):
+        # type: (Callable[..., None]) -> Callable[..., None]
+        """
+        Mark a method as a metadata entrypoint.
+
+        This decorator provides automatic no-op behavior in case metadata collection is disabled on the Agent.
+        """
+
+        @functools.wraps(method)
+        def entrypoint(self, *args, **kwargs):
+            # type: (AgentCheck, *Any, **Any) -> None
+            if not self.is_metadata_collection_enabled():
+                return
+
+            # NOTE: error handling still at the discretion of the wrapped method.
+            method(self, *args, **kwargs)
+
+        return entrypoint
+
     def set_external_tags(self, external_tags):
         # type: (Sequence[ExternalTagType]) -> None
         # Example of external_tags format

datadog_checks_base/tests/test_metadata.py

@@ -5,6 +5,7 @@
 import logging
 import re
 from collections import OrderedDict
+from typing import Any
 
 import mock
 import pytest
@@ -15,6 +16,7 @@
 pytestmark = pytest.mark.metadata
 
 SET_CHECK_METADATA_METHOD = 'datadog_checks.base.stubs.datadog_agent.set_check_metadata'
+GET_CONFIG_METHOD = 'datadog_checks.base.stubs.datadog_agent.get_config'
 
 # The order is used to derive the display name for the regex tests
 NON_STANDARD_VERSIONS = OrderedDict()
@@ -499,3 +501,41 @@ def test_blacklist_user_override(self):
             assert data.pop('is_set', None) is True
             assert data.pop('value', None) == 'bar'
             assert not data
+
+
+class TestMetadataEntrypoint:
+    def test_no_op_if_collection_disabled(self):
+        # type: () -> None
+
+        class MyCheck(AgentCheck):
+            @AgentCheck.metadata_entrypoint
+            def process_metadata(self, message):
+                # type: (str) -> None
+                self.set_metadata('my_message', message)
+
+            def check(self, instance):
+                # type: (Any) -> None
+                self.process_metadata(message='Hello, world')
+
+        with mock.patch(GET_CONFIG_METHOD, return_value=False):
+            check = MyCheck('test', {}, [{}])
+
+            with mock.patch(SET_CHECK_METADATA_METHOD) as m:
+                check.check({})
+                m.assert_not_called()
+
+    def test_exceptions_pass_through(self):
+        # type: (Any) -> None
+        class MyCheck(AgentCheck):
+            @AgentCheck.metadata_entrypoint
+            def process_metadata(self):
+                # type: () -> None
+                raise RuntimeError('Something went wrong')
+
+            def check(self, instance):
+                # type: (Any) -> None
+                self.process_metadata()
+
+        check = MyCheck('test', {}, [{}])
+        with pytest.raises(RuntimeError):
+            check.check({})