From bbda21f1ecf8d83ae41623b5f1347784817c00f9 Mon Sep 17 00:00:00 2001
From: Chris Rossi <chris@archimedeanco.com>
Date: Thu, 3 Sep 2020 15:03:06 -0400
Subject: [PATCH] fix: fix bug when setting naive datetime on
 `DateTimeProperty` with timezone

Fixes #517
---
 google/cloud/ndb/model.py |  2 ++
 tests/system/test_crud.py | 59 +++++++++++++++++++++++++++++++++++++++
 tests/unit/test_model.py  |  8 ++++++
 3 files changed, 69 insertions(+)

diff --git a/google/cloud/ndb/model.py b/google/cloud/ndb/model.py
index 302f96fa..03149486 100644
--- a/google/cloud/ndb/model.py
+++ b/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:
diff --git a/tests/system/test_crud.py b/tests/system/test_crud.py
index 4a4117ba..67a174c0 100644
--- a/tests/system/test_crud.py
+++ b/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):
diff --git a/tests/unit/test_model.py b/tests/unit/test_model.py
index 321dd544..3fa05d97 100644
--- a/tests/unit/test_model.py
+++ b/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")