fix: fix bug when setting naive datetime on DateTimeProperty with timezone (#534)

Fixes #517

Author: Chris Rossi

google/cloud/ndb/model.py

@@ -3747,6 +3747,8 @@ def _from_base_type(self, value):
             value = datetime.datetime.fromtimestamp(seconds, pytz.utc)
 
         if self._tzinfo is not None:
+            if value.tzinfo is None:
+                value = value.replace(tzinfo=pytz.utc)
             return value.astimezone(self._tzinfo)
 
         elif value.tzinfo is not None:

tests/system/test_crud.py

@@ -18,6 +18,7 @@
 import datetime
 import os
 import pickle
+import pytz
 import random
 import threading
 import zlib
@@ -40,6 +41,11 @@
 USE_REDIS_CACHE = bool(os.environ.get("REDIS_CACHE_URL"))
 
 
+def _assert_contemporaneous(timestamp1, timestamp2, delta_margin=2):
+    delta_margin = datetime.timedelta(seconds=delta_margin)
+    assert delta_margin > abs(timestamp1 - timestamp2)
+
+
 @pytest.mark.usefixtures("client_context")
 def test_retrieve_entity(ds_entity):
     entity_id = test_utils.system.unique_resource_id()
@@ -1043,6 +1049,59 @@ class SomeKind(ndb.Model):
     assert isinstance(retrieved.updated_at, datetime.datetime)
 
 
+@pytest.mark.usefixtures("client_context")
+def test_insert_autonow_property_with_tz(dispose_of):
+    """Regression test for #517
+
+    https://github.com/googleapis/python-ndb/issues/517
+    """
+
+    class SomeKind(ndb.Model):
+        created_at = ndb.DateTimeProperty(auto_now_add=True, tzinfo=pytz.utc)
+        updated_at = ndb.DateTimeProperty(auto_now=True, tzinfo=pytz.utc)
+
+    now = datetime.datetime.now(pytz.utc)
+    entity = SomeKind()
+    key = entity.put()
+    dispose_of(key._key)
+
+    _assert_contemporaneous(entity.created_at, now)
+    _assert_contemporaneous(entity.updated_at, now)
+
+    retrieved = key.get()
+
+    _assert_contemporaneous(retrieved.created_at, now)
+    _assert_contemporaneous(retrieved.updated_at, now)
+
+
+@pytest.mark.usefixtures("client_context")
+def test_insert_datetime_property_with_tz(dispose_of):
+    """Regression test for #517
+
+    https://github.com/googleapis/python-ndb/issues/517
+    """
+
+    class SomeKind(ndb.Model):
+        alarm1 = ndb.DateTimeProperty(tzinfo=pytz.utc)
+        alarm2 = ndb.DateTimeProperty(tzinfo=pytz.utc)
+
+    now = datetime.datetime.now(pytz.utc)
+    entity = SomeKind(
+        alarm1=now,
+        alarm2=datetime.datetime.utcnow(),  # naive
+    )
+    key = entity.put()
+    dispose_of(key._key)
+
+    _assert_contemporaneous(entity.alarm1, now)
+    _assert_contemporaneous(entity.alarm2, now)
+
+    retrieved = key.get()
+
+    _assert_contemporaneous(retrieved.alarm1, now)
+    _assert_contemporaneous(retrieved.alarm2, now)
+
+
 @pytest.mark.usefixtures("client_context")
 def test_insert_nested_autonow_property(dispose_of):
     class OtherKind(ndb.Model):

tests/unit/test_model.py

@@ -2784,6 +2784,14 @@ def test__from_base_type_convert_timezone():
             2010, 5, 11, 20, tzinfo=timezone(-4)
         )
 
+    @staticmethod
+    def test__from_base_type_naive_with_timezone():
+        prop = model.DateTimeProperty(name="dt_val", tzinfo=timezone(-4))
+        value = datetime.datetime(2010, 5, 12)
+        assert prop._from_base_type(value) == datetime.datetime(
+            2010, 5, 11, 20, tzinfo=timezone(-4)
+        )
+
     @staticmethod
     def test__from_base_type_int():
         prop = model.DateTimeProperty(name="dt_val")