fix misshandling mau reaping where reserved users are present

1 files changed, 49 insertions, 0 deletions

diff --git a/tests/storage/test_monthly_active_users.py b/tests/storage/test_monthly_active_users.py
index 1494650d10..62aa2bfed7 100644
--- a/tests/storage/test_monthly_active_users.py
+++ b/tests/storage/test_monthly_active_users.py
@@ -158,6 +158,54 @@ class MonthlyActiveUsersTestCase(unittest.HomeserverTestCase):
         count = self.store.get_monthly_active_count()
         self.assertEquals(self.get_success(count), 0)
 
+    def test_reap_monthly_active_users_reserved_users(self):
+        """ Tests that reaping correctly handles reaping where reserved users are
+        present"""
+
+        self.hs.config.max_mau_value = 5
+        initial_users = 5
+        reserved_user_number = initial_users - 1
+        threepids = []
+        for i in range(initial_users):
+            user = "@user%d:server" % i
+            email = "user%d@example.com" % i
+            self.store.upsert_monthly_active_user(user)
+            threepids.append({"medium": "email", "address": email},)
+            # Need to ensure that the most recent entries in the
+            # monthly_active_users table are reserved
+            now = int(self.hs.get_clock().time_msec())
+            if i != 0:
+                self.store.register_user(user_id=user, password_hash=None)
+                self.pump()
+                self.store.user_add_threepid(user, "email", email, now, now)
+
+        self.hs.config.mau_limits_reserved_threepids = threepids
+        self.store.runInteraction(
+            "initialise", self.store._initialise_reserved_users, threepids
+        )
+        self.pump()
+        count = self.store.get_monthly_active_count()
+        self.assertTrue(self.get_success(count), initial_users)
+
+        count = self.store.get_registered_reserved_users_count()
+        self.assertEquals(self.get_success(count), reserved_user_number)
+
+        self.store.reap_monthly_active_users()
+        self.pump()
+        count = self.store.get_monthly_active_count()
+        self.assertEquals(
+            self.get_success(count), self.hs.config.max_mau_value
+        )
+
+    def _populate_reserved_users(self, user, email):
+        """Helper function that registers userpopulates reserved users"""
+        # Test reserved registed users
+        self.store.register_user(user_id=user, password_hash=None)
+        self.pump()
+
+        now = int(self.hs.get_clock().time_msec())
+        self.store.user_add_threepid(user, "email", email, now, now)
+
     def test_populate_monthly_users_is_guest(self):
         # Test that guest users are not added to mau list
         user_id = "@user_id:host"
@@ -198,6 +246,7 @@ class MonthlyActiveUsersTestCase(unittest.HomeserverTestCase):
 
         user1 = "@user1:example.com"
         user2 = "@user2:example.com"
+
         user1_email = "user1@example.com"
         user2_email = "user2@example.com"
         threepids = [