Migrate the user directory initial population to a background task (#4864)

3 files changed, 79 insertions, 45 deletions

diff --git a/tests/handlers/test_user_directory.py b/tests/handlers/test_user_directory.py
index 114807efc1..aefe11ac28 100644
--- a/tests/handlers/test_user_directory.py
+++ b/tests/handlers/test_user_directory.py
@@ -163,9 +163,7 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
     def get_users_in_public_rooms(self):
         r = self.get_success(
             self.store._simple_select_list(
-                "users_in_public_rooms",
-                None,
-                ("user_id", "room_id"),
+                "users_in_public_rooms", None, ("user_id", "room_id")
             )
         )
         retval = []
@@ -182,6 +180,53 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
             )
         )
 
+    def _add_background_updates(self):
+        """
+        Add the background updates we need to run.
+        """
+        # Ugh, have to reset this flag
+        self.store._all_done = False
+
+        self.get_success(
+            self.store._simple_insert(
+                "background_updates",
+                {
+                    "update_name": "populate_user_directory_createtables",
+                    "progress_json": "{}",
+                },
+            )
+        )
+        self.get_success(
+            self.store._simple_insert(
+                "background_updates",
+                {
+                    "update_name": "populate_user_directory_process_rooms",
+                    "progress_json": "{}",
+                    "depends_on": "populate_user_directory_createtables",
+                },
+            )
+        )
+        self.get_success(
+            self.store._simple_insert(
+                "background_updates",
+                {
+                    "update_name": "populate_user_directory_process_users",
+                    "progress_json": "{}",
+                    "depends_on": "populate_user_directory_process_rooms",
+                },
+            )
+        )
+        self.get_success(
+            self.store._simple_insert(
+                "background_updates",
+                {
+                    "update_name": "populate_user_directory_cleanup",
+                    "progress_json": "{}",
+                    "depends_on": "populate_user_directory_process_users",
+                },
+            )
+        )
+
     def test_initial(self):
         """
         The user directory's initial handler correctly updates the search tables.
@@ -211,26 +256,17 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
         self.assertEqual(shares_private, [])
         self.assertEqual(public_users, [])
 
-        # Reset the handled users caches
-        self.handler.initially_handled_users = set()
+        # Do the initial population of the user directory via the background update
+        self._add_background_updates()
 
-        # Do the initial population
-        d = self.handler._do_initial_spam()
-
-        # This takes a while, so pump it a bunch of times to get through the
-        # sleep delays
-        for i in range(10):
-            self.pump(1)
-
-        self.get_success(d)
+        while not self.get_success(self.store.has_completed_background_updates()):
+            self.get_success(self.store.do_next_background_update(100), by=0.1)
 
         shares_private = self.get_users_who_share_private_rooms()
         public_users = self.get_users_in_public_rooms()
 
         # User 1 and User 2 are in the same public room
-        self.assertEqual(
-            set(public_users), set([(u1, room), (u2, room)])
-        )
+        self.assertEqual(set(public_users), set([(u1, room), (u2, room)]))
 
         # User 1 and User 3 share private rooms
         self.assertEqual(
@@ -238,7 +274,7 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
             set([(u1, u3, private_room), (u3, u1, private_room)]),
         )
 
-    def test_search_all_users(self):
+    def test_initial_share_all_users(self):
         """
         Search all users = True means that a user does not have to share a
         private room with the searching user or be in a public room to be search
@@ -248,33 +284,36 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
         self.hs.config.user_directory_search_all_users = True
 
         u1 = self.register_user("user1", "pass")
-        u1_token = self.login(u1, "pass")
-        u2 = self.register_user("user2", "pass")
-        u2_token = self.login(u2, "pass")
+        self.register_user("user2", "pass")
         u3 = self.register_user("user3", "pass")
 
-        # User 1 and User 2 join a room. User 3 never does.
-        room = self.helper.create_room_as(u1, is_public=True, tok=u1_token)
-        self.helper.invite(room, src=u1, targ=u2, tok=u1_token)
-        self.helper.join(room, user=u2, tok=u2_token)
-
+        # Wipe the user dir
         self.get_success(self.store.update_user_directory_stream_pos(None))
         self.get_success(self.store.delete_all_from_user_dir())
 
-        # Reset the handled users caches
-        self.handler.initially_handled_users = set()
+        # Do the initial population of the user directory via the background update
+        self._add_background_updates()
 
-        # Do the initial population
-        d = self.handler._do_initial_spam()
+        while not self.get_success(self.store.has_completed_background_updates()):
+            self.get_success(self.store.do_next_background_update(100), by=0.1)
 
-        # This takes a while, so pump it a bunch of times to get through the
-        # sleep delays
-        for i in range(10):
-            self.pump(1)
+        shares_private = self.get_users_who_share_private_rooms()
+        public_users = self.get_users_in_public_rooms()
 
-        self.get_success(d)
+        # No users share rooms
+        self.assertEqual(public_users, [])
+        self.assertEqual(self._compress_shared(shares_private), set([]))
 
         # Despite not sharing a room, search_all_users means we get a search
         # result.
         s = self.get_success(self.handler.search_users(u1, u3, 10))
         self.assertEqual(len(s["results"]), 1)
+
+        # We can find the other two users
+        s = self.get_success(self.handler.search_users(u1, "user", 10))
+        self.assertEqual(len(s["results"]), 2)
+
+        # Registering a user and then searching for them works.
+        u4 = self.register_user("user4", "pass")
+        s = self.get_success(self.handler.search_users(u1, u4, 10))
+        self.assertEqual(len(s["results"]), 1)
diff --git a/tests/storage/test_user_directory.py b/tests/storage/test_user_directory.py
index 512d76e7a3..fd3361404f 100644
--- a/tests/storage/test_user_directory.py
+++ b/tests/storage/test_user_directory.py
@@ -16,7 +16,6 @@
 from twisted.internet import defer
 
 from synapse.storage import UserDirectoryStore
-from synapse.storage.roommember import ProfileInfo
 
 from tests import unittest
 from tests.utils import setup_test_homeserver
@@ -34,13 +33,9 @@ class UserDirectoryStoreTestCase(unittest.TestCase):
 
         # alice and bob are both in !room_id. bobby is not but shares
         # a homeserver with alice.
-        yield self.store.add_profiles_to_user_dir(
-            {
-                ALICE: ProfileInfo(None, "alice"),
-                BOB: ProfileInfo(None, "bob"),
-                BOBBY: ProfileInfo(None, "bobby"),
-            },
-        )
+        yield self.store.update_profile_in_user_dir(ALICE, "alice", None)
+        yield self.store.update_profile_in_user_dir(BOB, "bob", None)
+        yield self.store.update_profile_in_user_dir(BOBBY, "bobby", None)
         yield self.store.add_users_in_public_rooms(
             "!room:id", (ALICE, BOB)
         )
diff --git a/tests/unittest.py b/tests/unittest.py
index ef31321bc8..7772a47078 100644
--- a/tests/unittest.py
+++ b/tests/unittest.py
@@ -330,10 +330,10 @@ class HomeserverTestCase(TestCase):
         """
         self.reactor.pump([by] * 100)
 
-    def get_success(self, d):
+    def get_success(self, d, by=0.0):
         if not isinstance(d, Deferred):
             return d
-        self.pump()
+        self.pump(by=by)
         return self.successResultOf(d)
 
     def register_user(self, username, password, admin=False):