bug: properly set expiration field.

Closes #1231
mozilla-services · May 14, 2018 · f764634 · f764634
1 parent 77dab57
commit f764634
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 14 deletions.
diff --git a/autopush/db.py b/autopush/db.py
@@ -538,25 +538,39 @@ def __init__(self, tablename, boto_resource=None,
     def table_status(self):
         return self.table.table_status
 
+    @track_provisioned
+    def touch_record(self, uaid, ttl=None):
+        # type: (str, int) -> bool
+        if ttl is None:
+            ttl = self._max_ttl
+        self.table.update_item(
+            Key={
+                'uaid': hasher(uaid),
+                'chidmessageid': ' ',
+            },
+            # `list_append(chids, :channel_id)` errors if chids doesn't exist
+            UpdateExpression='SET expiry = :expiry',
+            ExpressionAttributeValues={
+                ":expiry": _expiry(ttl)
+            },
+        )
+
     @track_provisioned
     def register_channel(self, uaid, channel_id, ttl=None):
         # type: (str, str, int) -> bool
         """Register a channel for a given uaid"""
-        # Generate our update expression
-        if ttl is None:
-            ttl = self._max_ttl
-        expr_values = {
-            ":channel_id": set([normalize_id(channel_id)]),
-            ":expiry": _expiry(ttl)
-        }
+        # Create or Add the channelID to the list of channels
         self.table.update_item(
             Key={
                 'uaid': hasher(uaid),
                 'chidmessageid': ' ',
             },
-            UpdateExpression='ADD chids :channel_id, expiry :expiry',
-            ExpressionAttributeValues=expr_values,
+            UpdateExpression='ADD chids :channel_id',
+            ExpressionAttributeValues={
+                ":channel_id": set([normalize_id(channel_id)])
+            },
         )
+        self.touch_record(uaid, ttl)
         return True
 
     @track_provisioned
@@ -579,6 +593,7 @@ def unregister_channel(self, uaid, channel_id, **kwargs):
         chids = response.get('Attributes', {}).get('chids', {})
         if chids:
             try:
+                self.touch_record(uaid)
                 return chid in chids
             except (TypeError, AttributeError):  # pragma: nocover
                 pass
@@ -604,6 +619,7 @@ def all_channels(self, uaid):
             return False, set([])
         if 'Item' not in result:
             return False, set([])
+        self.touch_record(uaid)
         return True, result['Item'].get("chids", set([]))
 
     @track_provisioned
@@ -694,7 +710,6 @@ def fetch_messages(
             # Ensure we return an int, as boto2 can return Decimals
             if results[0].get("current_timestamp"):
                 last_position = int(results[0]["current_timestamp"])
-
         return last_position, [
             WebPushNotification.from_message_table(uaid, x)
             for x in results[1:]

diff --git a/autopush/tests/test_integration.py b/autopush/tests/test_integration.py
@@ -615,7 +615,7 @@ def _ws_url(self):
         return self.conn.conf.ws_url
 
     @inlineCallbacks
-    def test_hello_only_has_three_calls(self):
+    def test_hello_calls(self):
         db.TRACK_DB_CALLS = True
         client = Client(self._ws_url)
         yield client.connect()

diff --git a/autopush/tests/test_rs_integration.py b/autopush/tests/test_rs_integration.py
@@ -217,7 +217,7 @@ def test_no_rotation(self):
         yield self.shut_down(client)
 
     @inlineCallbacks
-    def test_hello_only_has_three_calls(self):
+    def test_hello_calls(self):
         db.TRACK_DB_CALLS = True
         client = Client(self._ws_url)
         yield client.connect()
@@ -234,7 +234,8 @@ def test_hello_only_has_three_calls(self):
         yield client.wait_for(lambda: len(db.DB_CALLS) == 4)
         assert db.DB_CALLS == ['register_user', 'register_user',
                                'fetch_messages',
-                               'fetch_timestamp_messages']
+                               'fetch_timestamp_messages',
+                               'touch_record']
         db.DB_CALLS = []
         db.TRACK_DB_CALLS = False
 

diff --git a/autopush/webpush_server.py b/autopush/webpush_server.py
@@ -490,7 +490,7 @@ def _check_storage(self, command):
             )
         messages = [WebPushMessage.from_WebPushNotification(m)
                     for m in messages]
-
+        message.touch_record(command.uaid.hex)
         # If we're out of messages, timestamp is set to None, so we return
         # the last timestamp supplied
         if not timestamp: