Fix bug (GoogleCloudPlatform#1709)

Ace Nassri · web-flow · commit eac8dafcf8bf · 2018-09-28T18:06:28.000-07:00
Change-Id: Ic1fece4ad23523038702068a87b5ec2f1857f18f
diff --git a/functions/tips/main.py b/functions/tips/main.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 # [START functions_tips_infinite_retries]
-from datetime import datetime
+from datetime import datetime, timezone
 # [END functions_tips_infinite_retries]
 
 # [START functions_tips_gcp_apis]
@@ -150,19 +150,20 @@ def avoid_infinite_retries(data, context):
         None; output is written to Stackdriver Logging
     """
 
-    timestamp = data.timestamp
+    timestamp = context.timestamp
 
     event_time = parser.parse(timestamp)
-    event_age = (datetime.now() - event_time).total_seconds() * 1000
+    event_age = (datetime.now(timezone.utc) - event_time).total_seconds()
+    event_age_ms = event_age * 1000
 
     # Ignore events that are too old
     max_age_ms = 10000
-    if event_age > max_age_ms:
-        print('Dropped {} (age {}ms)'.format(context.event_id, event_age))
+    if event_age_ms > max_age_ms:
+        print('Dropped {} (age {}ms)'.format(context.event_id, event_age_ms))
         return 'Timeout'
 
     # Do what the function is supposed to do
-    print('Processed {} (age {}ms)'.format(context.event_id, event_age))
+    print('Processed {} (age {}ms)'.format(context.event_id, event_age_ms))
     return
 # [END functions_tips_infinite_retries]
 
diff --git a/functions/tips/main_test.py b/functions/tips/main_test.py
@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import datetime
+from collections import UserDict
+from datetime import datetime, timedelta, timezone
 from unittest.mock import MagicMock, Mock, patch
 
 import flask
@@ -52,23 +53,26 @@ def test_connection_pooling_404(app):
 
 
 def test_avoid_infinite_retries(capsys):
-    now = datetime.datetime.now()
+    now = datetime.now(timezone.utc)
 
-    with patch('main.datetime', wraps=datetime.datetime) as datetime_mock:
+    with patch('main.datetime', wraps=datetime) as datetime_mock:
         datetime_mock.now = Mock(return_value=now)
-        old_event = Mock(
-            timestamp=(now - datetime.timedelta(seconds=15)).isoformat())
-        young_event = Mock(
-            timestamp=(now - datetime.timedelta(seconds=5)).isoformat())
-        context = Mock(event_id='fake_event_id')
 
-        main.avoid_infinite_retries(old_event, context)
+        old_context = UserDict()
+        old_context.timestamp = (now - timedelta(seconds=15)).isoformat()
+        old_context.event_id = 'old_event_id'
+
+        young_context = UserDict()
+        young_context.timestamp = (now - timedelta(seconds=5)).isoformat()
+        young_context.event_id = 'young_event_id'
+
+        main.avoid_infinite_retries(None, old_context)
         out, _ = capsys.readouterr()
-        assert 'Dropped {} (age 15000.0ms)'.format(context.event_id) in out
+        assert f"Dropped {old_context.event_id} (age 15000.0ms)" in out
 
-        main.avoid_infinite_retries(young_event, context)
+        main.avoid_infinite_retries(None, young_context)
         out, _ = capsys.readouterr()
-        assert 'Processed {} (age 5000.0ms)'.format(context.event_id) in out
+        assert f"Processed {young_context.event_id} (age 5000.0ms)" in out
 
 
 def test_retry_or_not():
