Even better

DataDog · Jun 22, 2018 · 0e74b60 · 0e74b60
1 parent 767f162
commit 0e74b60
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 12 deletions.
diff --git a/datadog_checks_base/datadog_checks/checks/prometheus/base_check.py b/datadog_checks_base/datadog_checks/checks/prometheus/base_check.py
@@ -165,7 +165,7 @@ def get_scraper(self, instance):
         scraper.ssl_private_key = instance.get("ssl_private_key", default_instance.get("ssl_private_key", None))
         scraper.ssl_ca_cert = instance.get("ssl_ca_cert", default_instance.get("ssl_ca_cert", None))
 
-        scraper.prometheus_timeout = scraper._extract_prometheus_timeout(instance, default_instance.get("prometheus_timeout", 10))
+        scraper.set_prometheus_timeout(instance, default_instance.get("prometheus_timeout", 10))
 
         self.scrapers_map[endpoint] = scraper
 

diff --git a/datadog_checks_base/datadog_checks/checks/prometheus/mixins.py b/datadog_checks_base/datadog_checks/checks/prometheus/mixins.py
@@ -356,11 +356,6 @@ def process(self, endpoint, **kwargs):
         if instance:
             kwargs['custom_tags'] = instance.get('tags', [])
 
-        # Used in a regular PrometheusCheck, a prometheus_timeout can be passed here as an argument
-        timeout = kwargs.get('prometheus_timeout')
-        if timeout:
-            self.prometheus_timeout = timeout
-
         for metric in self.scrape_metrics(endpoint):
             self.process_metric(metric, **kwargs)
 
@@ -631,11 +626,12 @@ def _submit_gauges_from_histogram(self, name, metric, send_histograms_buckets=Tr
     def _is_value_valid(self, val):
         return not (isnan(val) or isinf(val))
 
-    def _extract_prometheus_timeout(self, instance, default_value=10):
+    def set_prometheus_timeout(self, instance, default_value=10):
+        """ extract `prometheus_timeout` directly from the instance configuration """
+        self.prometheus_timeout = default_value
         timeout = instance.get('prometheus_timeout')
         if timeout is not None:
             if timeout <= 0:
                 self.log.debug("Prometheus integration timeout is incorrect, defaulting to {}".format(default_value))
-                return(default_value)
-            return(timeout)
-        return(default_value)
+            else:
+                self.prometheus_timeout = timeout
diff --git a/kubernetes_state/datadog_checks/kubernetes_state/kubernetes_state.py b/kubernetes_state/datadog_checks/kubernetes_state/kubernetes_state.py
@@ -177,7 +177,7 @@ def check(self, instance):
         if endpoint is None:
             raise CheckException("Unable to find kube_state_url in config file.")
 
-        timeout = self._extract_prometheus_timeout(instance)
+        self.set_prometheus_timeout(instance)
 
         if 'labels_mapper' in instance:
             if isinstance(instance['labels_mapper'], dict):
@@ -200,7 +200,7 @@ def check(self, instance):
         self.job_succeeded_count = defaultdict(int)
         self.job_failed_count = defaultdict(int)
 
-        self.process(endpoint, send_histograms_buckets=send_buckets, instance=instance, prometheus_timeout=timeout)
+        self.process(endpoint, send_histograms_buckets=send_buckets, instance=instance)
 
         for job_tags, job_count in self.job_succeeded_count.iteritems():
             self.monotonic_count(self.NAMESPACE + '.job.succeeded', job_count, list(job_tags))