Safely handle datetime constraint values with and without timezone in…

…fo (#324)
Unleash · Nov 15, 2024 · ced8668 · ced8668
1 parent b815f6d
commit ced8668
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 6 deletions.
diff --git a/UnleashClient/constraints/Constraint.py b/UnleashClient/constraints/Constraint.py
@@ -1,5 +1,5 @@
 # pylint: disable=invalid-name, too-few-public-methods, use-a-generator
-from datetime import datetime
+from datetime import datetime, timezone
 from enum import Enum
 from typing import Any, Optional, Union
 
@@ -125,6 +125,11 @@ def check_numeric_operators(self, context_value: Union[float, int]) -> bool:
         return return_value
 
     def check_date_operators(self, context_value: Union[datetime, str]) -> bool:
+        if isinstance(context_value, datetime) and context_value.tzinfo is None:
+            raise ValueError(
+                "If context_value is a datetime object, it must be timezone (offset) aware."
+            )
+
         return_value = False
         parsing_exception = False
 
@@ -139,8 +144,16 @@ def check_date_operators(self, context_value: Union[datetime, str]) -> bool:
 
         try:
             parsed_date = parse(self.value)
+
+            # If parsed date is timezone-naive, assume it is UTC
+            if parsed_date.tzinfo is None:
+                parsed_date = parsed_date.replace(tzinfo=timezone.utc)
+
             if isinstance(context_value, str):
                 context_date = parse(context_value)
+                # If parsed date is timezone-naive, assume it is UTC
+                if context_date.tzinfo is None:
+                    context_date = context_date.replace(tzinfo=timezone.utc)
             else:
                 context_date = context_value
         except DateUtilParserError:
@@ -197,7 +210,8 @@ def apply(self, context: dict = None) -> bool:
 
             # Set currentTime if not specified
             if self.context_name == "currentTime" and not context_value:
-                context_value = datetime.now()
+                # Use the current system time in the local timezone (tz-aware)
+                context_value = datetime.now(timezone.utc).astimezone()
 
             if context_value is not None:
                 if self.operator in [

diff --git a/tests/unit_tests/test_constraints.py b/tests/unit_tests/test_constraints.py
@@ -1,4 +1,4 @@
-from datetime import datetime
+from datetime import datetime, timedelta
 
 import pytest
 import pytz
@@ -183,15 +183,61 @@ def test_constraints_DATE_BEFORE():
     assert constraint.apply({"currentTime": datetime(2022, 1, 21, tzinfo=pytz.UTC)})
 
 
-def test_constraints_default():
-    constraint = Constraint(constraint_dict=mock_constraints.CONSTRAINT_DATE_BEFORE)
+def test_constraints_DATE_AFTER_default():
+    constraint = Constraint(
+        constraint_dict={
+            **mock_constraints.CONSTRAINT_DATE_AFTER,
+            "value": (datetime.now(pytz.UTC) - timedelta(days=1)).isoformat(),
+        }
+    )
+
+    assert constraint.apply({})
+
+    constraint = Constraint(
+        constraint_dict={
+            **mock_constraints.CONSTRAINT_DATE_AFTER,
+            "value": (datetime.now(pytz.UTC) + timedelta(days=1)).isoformat(),
+        }
+    )
+
+    assert not constraint.apply({})
+
+
+def test_constraints_DATE_BEFORE_default():
+    constraint = Constraint(
+        constraint_dict={
+            **mock_constraints.CONSTRAINT_DATE_BEFORE,
+            "value": (datetime.now(pytz.UTC) + timedelta(days=1)).isoformat(),
+        }
+    )
+
+    assert constraint.apply({})
+
+    constraint = Constraint(
+        constraint_dict={
+            **mock_constraints.CONSTRAINT_DATE_BEFORE,
+            "value": (datetime.now(pytz.UTC) - timedelta(days=1)).isoformat(),
+        }
+    )
 
     assert not constraint.apply({})
 
 
+def test_constraints_tz_naive():
+    constraint = Constraint(constraint_dict=mock_constraints.CONSTRAINT_DATE_TZ_NAIVE)
+
+    assert constraint.apply(
+        {"currentTime": datetime(2022, 1, 22, 0, 10, tzinfo=pytz.UTC)}
+    )
+    assert not constraint.apply({"currentTime": datetime(2022, 1, 22, tzinfo=pytz.UTC)})
+    assert not constraint.apply(
+        {"currentTime": datetime(2022, 1, 21, 23, 50, tzinfo=pytz.UTC)}
+    )
+
+
 def test_constraints_date_error():
     constraint = Constraint(constraint_dict=mock_constraints.CONSTRAINT_DATE_ERROR)
-    assert not constraint.apply({"currentTime": datetime(2022, 1, 23)})
+    assert not constraint.apply({"currentTime": datetime(2022, 1, 23, tzinfo=pytz.UTC)})
 
 
 def test_constraints_SEMVER_EQ():

diff --git a/tests/utilities/mocks/mock_constraints.py b/tests/utilities/mocks/mock_constraints.py
@@ -153,6 +153,13 @@
     "inverted": False,
 }
 
+CONSTRAINT_DATE_TZ_NAIVE = {
+    "contextName": "currentTime",
+    "operator": "DATE_AFTER",
+    "value": "2022-01-22T00:00:00.000",
+    "inverted": False,
+}
+
 
 CONSTRAINT_SEMVER_EQ = {
     "contextName": "customField",