fix py2

DataDog · Dec 8, 2021 · a4ab9b1 · a4ab9b1
1 parent 0ccfac7
commit a4ab9b1
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/datadog_checks_base/datadog_checks/base/utils/tracking.py b/datadog_checks_base/datadog_checks/base/utils/tracking.py
@@ -4,7 +4,8 @@
 # Licensed under a 3-clause BSD style license (see LICENSE)
 
 import os
-import time
+
+from datadog_checks.base.utils.time import get_timestamp
 
 required_attrs = [
     'name',
@@ -43,7 +44,7 @@ def wrapper(self, *args, **kwargs):
             if os.getenv('DD_DISABLE_TRACKED_METHOD') == "true":
                 return function(self, *args, **kwargs)
 
-            start_time = time.time()
+            start_time = get_timestamp()
 
             try:
                 check = agent_check_getter(self) if agent_check_getter else self
@@ -69,7 +70,7 @@ def wrapper(self, *args, **kwargs):
             try:
                 result = function(self, *args, **kwargs)
 
-                elapsed_ms = (time.time() - start_time) * 1000
+                elapsed_ms = (get_timestamp() - start_time) * 1000
                 check.histogram("dd.{}.operation.time".format(check_name), elapsed_ms, **stats_kwargs)
 
                 check.log.debug("[%s.%s] operation completed in %s ms", check_name, function.__name__, elapsed_ms)

diff --git a/datadog_checks_base/tests/base/utils/test_tracking.py b/datadog_checks_base/tests/base/utils/test_tracking.py
@@ -27,6 +27,10 @@ def debug_stats_kwargs(self):
 EXPECTED_RESULT = 5
 
 
+class MyException(Exception):
+    pass
+
+
 class TestJob:
     def __init__(self, check):
         self.check = check
@@ -51,7 +55,7 @@ def do_work_return_list(self):
 
     @tracked_method(agent_check_getter=agent_check_getter)
     def test_tracked_exception(self):
-        raise Exception("oops")
+        raise MyException("oops")
 
 
 @pytest.mark.parametrize(
@@ -89,5 +93,5 @@ def test_tracked_method(aggregator, debug_stats_kwargs, disable_tracking):
         aggregator.assert_metric(
             "dd.hello.operation.error",
             hostname=hostname,
-            tags=tags + ["operation:test_tracked_exception", "error:<class 'Exception'>"],
+            tags=tags + ["operation:test_tracked_exception", "error:<class 'test_tracking.MyException'>"],
         )