feat: port migrate user to Rust, remove Python calling

Closes #1206

autopush/tests/test_rs_integration.py

@@ -26,11 +26,17 @@
 import twisted.internet.base
 from cryptography.fernet import Fernet
 from twisted.internet.defer import inlineCallbacks, returnValue
+from twisted.internet.threads import deferToThread
 from twisted.trial import unittest
 from twisted.logger import globalLogPublisher
 
 import autopush.tests
 from autopush.config import AutopushConfig
+from autopush.db import (
+    get_month,
+    has_connected_this_month,
+    Message,
+)
 from autopush.logging import begin_or_register
 from autopush.main import EndpointApplication, RustConnectionApplication
 from autopush.utils import base64url_encode
@@ -733,6 +739,269 @@ def test_delete_saved_notification(self):
         assert result is None
         yield self.shut_down(client)
 
+    @inlineCallbacks
+    def test_webpush_monthly_rotation(self):
+        from autopush.db import make_rotating_tablename
+        client = yield self.quick_register()
+        yield client.disconnect()
+
+        # Move the client back one month to the past
+        last_month = make_rotating_tablename(
+            prefix=self.conn.conf.message_table.tablename, delta=-1)
+        lm_message = Message(last_month, boto_resource=self.conn.db.resource)
+        yield deferToThread(
+            self.conn.db.router.update_message_month,
+            client.uaid,
+            last_month,
+        )
+
+        # Verify the move
+        c = yield deferToThread(self.conn.db.router.get_uaid,
+                                client.uaid)
+        assert c["current_month"] == last_month
+
+        # Verify last_connect is current, then move that back
+        assert has_connected_this_month(c)
+        today = get_month(delta=-1)
+        last_connect = int("%s%s020001" % (today.year,
+                                           str(today.month).zfill(2)))
+
+        yield deferToThread(
+            self.conn.db.router._update_last_connect,
+            client.uaid,
+            last_connect)
+        c = yield deferToThread(self.conn.db.router.get_uaid,
+                                client.uaid)
+        assert has_connected_this_month(c) is False
+
+        # Move the clients channels back one month
+        exists, chans = yield deferToThread(
+            self.conn.db.message.all_channels,
+            client.uaid
+        )
+        assert exists is True
+        assert len(chans) == 1
+        yield deferToThread(
+            lm_message.save_channels,
+            client.uaid,
+            chans,
+        )
+
+        # Remove the channels entry entirely from this month
+        yield deferToThread(
+            self.conn.db.message.table.delete_item,
+            Key={'uaid': client.uaid, 'chidmessageid': ' '}
+        )
+
+        # Verify the channel is gone
+        exists, chans = yield deferToThread(
+            self.conn.db.message.all_channels,
+            client.uaid,
+        )
+        assert exists is False
+        assert len(chans) == 0
+
+        # Send in a notification, verify it landed in last months notification
+        # table
+        data = uuid.uuid4().hex
+        with self.legacy_endpoint():
+            yield client.send_notification(data=data)
+        ts, notifs = yield deferToThread(lm_message.fetch_timestamp_messages,
+                                         uuid.UUID(client.uaid),
+                                         " ")
+        assert len(notifs) == 1
+
+        # Connect the client, verify the migration
+        yield client.connect()
+        yield client.hello()
+
+        # Pull down the notification
+        result = yield client.get_notification()
+        chan = client.channels.keys()[0]
+        assert result is not None
+        assert chan == result["channelID"]
+
+        # Acknowledge the notification, which triggers the migration
+        yield client.ack(chan, result["version"])
+
+        # Wait up to 4 seconds for the table rotation to occur
+        start = time.time()
+        while time.time()-start < 4:
+            c = yield deferToThread(
+                self.conn.db.router.get_uaid,
+                client.uaid)
+            if c["current_month"] == self.conn.db.current_msg_month:
+                break
+            else:
+                yield deferToThread(time.sleep, 0.2)
+
+        # Verify the month update in the router table
+        c = yield deferToThread(
+            self.conn.db.router.get_uaid,
+            client.uaid)
+        assert c["current_month"] == self.conn.db.current_msg_month
+
+        # Verify the client moved last_connect
+        assert has_connected_this_month(c) is True
+
+        # Verify the channels were moved
+        exists, chans = yield deferToThread(
+            self.conn.db.message.all_channels,
+            client.uaid
+        )
+        assert exists is True
+        assert len(chans) == 1
+        yield self.shut_down(client)
+
+    @inlineCallbacks
+    def test_webpush_monthly_rotation_prior_record_exists(self):
+        from autopush.db import make_rotating_tablename
+        client = yield self.quick_register()
+        yield client.disconnect()
+
+        # Move the client back one month to the past
+        last_month = make_rotating_tablename(
+            prefix=self.conn.conf.message_table.tablename, delta=-1)
+        lm_message = Message(last_month,
+                             boto_resource=autopush.tests.boto_resource)
+        yield deferToThread(
+            self.conn.db.router.update_message_month,
+            client.uaid,
+            last_month,
+        )
+
+        # Verify the move
+        c = yield deferToThread(self.conn.db.router.get_uaid,
+                                client.uaid)
+        assert c["current_month"] == last_month
+
+        # Verify last_connect is current, then move that back
+        assert has_connected_this_month(c)
+        today = get_month(delta=-1)
+        yield deferToThread(
+            self.conn.db.router._update_last_connect,
+            client.uaid,
+            int("%s%s020001" % (today.year, str(today.month).zfill(2))),
+        )
+        c = yield deferToThread(self.conn.db.router.get_uaid, client.uaid)
+        assert has_connected_this_month(c) is False
+
+        # Move the clients channels back one month
+        exists, chans = yield deferToThread(
+            self.conn.db.message.all_channels,
+            client.uaid,
+        )
+        assert exists is True
+        assert len(chans) == 1
+        yield deferToThread(
+            lm_message.save_channels,
+            client.uaid,
+            chans,
+        )
+
+        # Send in a notification, verify it landed in last months notification
+        # table
+        data = uuid.uuid4().hex
+        with self.legacy_endpoint():
+            yield client.send_notification(data=data)
+        _, notifs = yield deferToThread(lm_message.fetch_timestamp_messages,
+                                        uuid.UUID(client.uaid),
+                                        " ")
+        assert len(notifs) == 1
+
+        # Connect the client, verify the migration
+        yield client.connect()
+        yield client.hello()
+
+        # Pull down the notification
+        result = yield client.get_notification()
+        chan = client.channels.keys()[0]
+        assert result is not None
+        assert chan == result["channelID"]
+
+        # Acknowledge the notification, which triggers the migration
+        yield client.ack(chan, result["version"])
+
+        # Wait up to 4 seconds for the table rotation to occur
+        start = time.time()
+        while time.time()-start < 4:
+            c = yield deferToThread(
+                self.conn.db.router.get_uaid,
+                client.uaid)
+            if c["current_month"] == self.conn.db.current_msg_month:
+                break
+            else:
+                yield deferToThread(time.sleep, 0.2)
+
+        # Verify the month update in the router table
+        c = yield deferToThread(self.conn.db.router.get_uaid, client.uaid)
+        assert c["current_month"] == self.conn.db.current_msg_month
+
+        # Verify the client moved last_connect
+        assert has_connected_this_month(c) is True
+
+        # Verify the channels were moved
+        exists, chans = yield deferToThread(
+            self.conn.db.message.all_channels,
+            client.uaid
+        )
+        assert exists is True
+        assert len(chans) == 1
+        yield self.shut_down(client)
+
+    @inlineCallbacks
+    def test_webpush_monthly_rotation_no_channels(self):
+        from autopush.db import make_rotating_tablename
+        client = Client("ws://localhost:{}/".format(self.connection_port))
+        yield client.connect()
+        yield client.hello()
+        yield client.disconnect()
+
+        # Move the client back one month to the past
+        last_month = make_rotating_tablename(
+            prefix=self.conn.conf.message_table.tablename, delta=-1)
+        yield deferToThread(
+            self.conn.db.router.update_message_month,
+            client.uaid,
+            last_month
+        )
+
+        # Verify the move
+        c = yield deferToThread(self.conn.db.router.get_uaid,
+                                client.uaid
+                                )
+        assert c["current_month"] == last_month
+
+        # Verify there's no channels
+        exists, chans = yield deferToThread(
+            self.conn.db.message.all_channels,
+            client.uaid,
+        )
+        assert exists is False
+        assert len(chans) == 0
+
+        # Connect the client, verify the migration
+        yield client.connect()
+        yield client.hello()
+
+        # Wait up to 2 seconds for the table rotation to occur
+        start = time.time()
+        while time.time()-start < 2:
+            c = yield deferToThread(
+                self.conn.db.router.get_uaid,
+                client.uaid,
+            )
+            if c["current_month"] == self.conn.db.current_msg_month:
+                break
+            else:
+                yield deferToThread(time.sleep, 0.2)
+
+        # Verify the month update in the router table
+        c = yield deferToThread(self.conn.db.router.get_uaid,
+                                client.uaid)
+        assert c["current_month"] == self.conn.db.current_msg_month
+        yield self.shut_down(client)
+
     @inlineCallbacks
     def test_with_key(self):
         private_key = ecdsa.SigningKey.generate(curve=ecdsa.NIST256p)