initial

1 files changed, 12 insertions, 0 deletions

diff --git a/tests/handlers/test_user_directory.py b/tests/handlers/test_user_directory.py
index a16a2dc67b..0e0ac0a48b 100644
--- a/tests/handlers/test_user_directory.py
+++ b/tests/handlers/test_user_directory.py
@@ -121,6 +121,7 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
         self.assertEqual(
             self._compress_shared(shares_private), set([(u1, u2, room), (u2, u1, room)])
         )
+        self.assertEqual(set(public_users), set([u1, u2]))
 
         # We get one search result when searching for user2 by user1.
         s = self.get_success(self.handler.search_users(u1, "user2", 10))
@@ -140,9 +141,11 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
         # Check we have removed the values.
         shares_public = self.get_users_who_share_public_rooms()
         shares_private = self.get_users_who_share_private_rooms()
+        public_users = self.get_users_in_public_rooms()
 
         self.assertEqual(shares_public, [])
         self.assertEqual(self._compress_shared(shares_private), set())
+        self.assertEqual(public_users, [u1])
 
         # User1 now gets no search results for any of the other users.
         s = self.get_success(self.handler.search_users(u1, "user2", 10))
@@ -160,6 +163,15 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
             r.add((i["user_id"], i["other_user_id"], i["room_id"]))
         return r
 
+    def get_users_in_public_rooms(self):
+        return self.get_success(
+            self.store._simple_select_list(
+                "users_in_public_rooms",
+                None,
+                ["user_id"],
+            )
+        )
+
     def get_users_who_share_public_rooms(self):
         return self.get_success(
             self.store._simple_select_list(