Fix time utilities (#6692)

* Fix time utilities

* address feedback - assert returned dt is same object

* address feedback - remove upstream test

* add comment

* remove comment

* address feedback

* address

Author: ofek

datadog_checks_base/datadog_checks/base/utils/time.py

@@ -6,68 +6,40 @@
 from datetime import datetime
 from time import time as epoch_offset
 
-import pytz
-from six import PY2
+from dateutil.tz import UTC
 
-UTC = pytz.utc
-EPOCH = datetime(1970, 1, 1, tzinfo=UTC)
+EPOCH = datetime.fromtimestamp(0, UTC)
 
 
-if PY2:
-
-    def get_timestamp(dt=None):
-        """
-        Returns the number of seconds since the Unix epoch.
-        If `dt` is not specified or `None`, the current time in UTC is assumed.
-        """
-        if dt is None:
-            return epoch_offset()
-
-        return (dt - EPOCH).total_seconds()
-
-
-else:
-
-    def get_timestamp(dt=None):
-        """
-        Returns the number of seconds since the Unix epoch.
-        If `dt` is not specified or `None`, the current time in UTC is assumed.
-        """
-        if dt is None:
-            # NOTE: Although the epoch is platform dependent, it appears to be the same
-            # for all platforms we've tested, therefore we use `time.time` for speed.
-            #
-            # Here is the test:
-            #
-            #     $ python -c "import time;print(tuple(time.gmtime(0)[:9]))"
-            #     (1970, 1, 1, 0, 0, 0, 3, 1, 0)
-            #
-            # If you can reproduce, add to the following list of tested platforms:
-            #
-            # - Windows
-            # - macOS
-            # - Ubuntu
-            # - Alpine
-            return epoch_offset()
+def get_timestamp(dt=None):
+    """
+    Returns the number of seconds since the Unix epoch.
+    If `dt` is not specified or `None`, the current time in UTC is assumed.
+    """
+    if dt is None:
+        # The precision is different between Python 2 and 3
+        return epoch_offset()
 
-        return normalize_datetime(dt).timestamp()
+    # TODO: when we drop support for Python 2 switch to:
+    # return ensure_aware_datetime(dt).timestamp()
+    return (ensure_aware_datetime(dt) - EPOCH).total_seconds()
 
 
-def get_aware_datetime(dt=None, tz=UTC):
+def get_current_datetime(tz=UTC):
     """
-    Returns an aware datetime object. If `dt` is not specified or `None`, the current time in UTC is assumed.
+    Returns an aware datetime object representing the current time. If `tz` is not specified, UTC is assumed.
     """
-    if dt is None:
-        return datetime.now(tz)
-    else:
-        return normalize_datetime(dt)
+    return datetime.now(tz)
 
 
-def normalize_datetime(dt, default_tz=UTC):
+def ensure_aware_datetime(dt, default_tz=UTC):
     """
     Ensures that the returned datetime object is not naive.
     """
     if dt.tzinfo is None:
         dt = dt.replace(tzinfo=default_tz)
 
     return dt
+
+
+__all__ = ['EPOCH', 'UTC', 'ensure_aware_datetime', 'get_current_datetime', 'get_timestamp']

datadog_checks_base/requirements.in

@@ -10,6 +10,7 @@ orjson==2.6.1; python_version > '3.0'
 prometheus-client==0.7.1
 protobuf==3.7.0
 pysocks==1.7.0
+python-dateutil==2.8.0
 pytz==2019.3
 pywin32==227; sys_platform == 'win32'
 pyyaml==5.3.1

datadog_checks_base/tests/test_time.py

@@ -0,0 +1,75 @@
+# (C) Datadog, Inc. 2020-present
+# All rights reserved
+# Licensed under a 3-clause BSD style license (see LICENSE)
+from datetime import datetime
+
+import pytest
+from dateutil import tz
+from six import PY2
+
+from datadog_checks.base.utils.time import EPOCH, UTC, ensure_aware_datetime, get_current_datetime, get_timestamp
+
+pytestmark = pytest.mark.time
+
+
+class TestNormalization:
+    def test_replace(self):
+        now = datetime.now()
+        assert now.tzinfo is None
+
+        assert ensure_aware_datetime(now).tzinfo is UTC
+
+    def test_utc(self):
+        nyc = tz.gettz('America/New_York')
+        now = datetime.now(nyc)
+
+        normalized = ensure_aware_datetime(now)
+
+        assert normalized is now
+        assert normalized.tzinfo is nyc
+
+
+class TestCurrentDatetime:
+    # What follows is a workaround for being unable to patch directly:
+    # TypeError: can't set attributes of built-in/extension type 'datetime.datetime'
+
+    def test_default_utc(self, mocker):
+        datetime_obj = mocker.patch('datadog_checks.base.utils.time.datetime')
+        datetime_obj.now = mocker.MagicMock()
+        dt = datetime(2020, 2, 2, tzinfo=UTC)
+        datetime_obj.now.return_value = dt
+        assert get_current_datetime() is dt
+        datetime_obj.now.assert_called_once_with(UTC)
+
+    def test_tz(self, mocker):
+        datetime_obj = mocker.patch('datadog_checks.base.utils.time.datetime')
+        datetime_obj.now = mocker.MagicMock()
+        nyc = tz.gettz('America/New_York')
+        dt = datetime(2020, 2, 2, tzinfo=nyc)
+        datetime_obj.now.return_value = dt
+        assert get_current_datetime(nyc) is dt
+        datetime_obj.now.assert_called_once_with(nyc)
+
+
+class TestTimestamp:
+    def test_type(self):
+        assert isinstance(get_timestamp(), float)
+
+    def test_default(self, mocker):
+        time_time = mocker.patch('datadog_checks.base.utils.time.epoch_offset')
+        get_timestamp()
+        time_time.assert_called_once()
+
+    def test_time_delta(self):
+        now = datetime.now()
+        expected = (now.replace(tzinfo=UTC) - EPOCH).total_seconds()
+
+        assert get_timestamp(now) == expected
+
+
+@pytest.mark.skipif(PY2, reason='Using Python 3 features')
+class TestConstants:
+    def test_epoch(self):
+        from datetime import timezone
+
+        assert EPOCH == datetime(1970, 1, 1, tzinfo=timezone.utc)