diff --git a/tests/handlers/test_user_directory.py b/tests/handlers/test_user_directory.py
index a16a2dc67b..114807efc1 100644
--- a/tests/handlers/test_user_directory.py
+++ b/tests/handlers/test_user_directory.py
@@ -114,13 +114,13 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
self.helper.join(room, user=u2, tok=u2_token)
# Check we have populated the database correctly.
- 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([(u1, u2, room), (u2, u1, room)])
)
+ self.assertEqual(public_users, [])
# We get one search result when searching for user2 by user1.
s = self.get_success(self.handler.search_users(u1, "user2", 10))
@@ -138,11 +138,11 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
self.helper.leave(room, user=u2, tok=u2_token)
# 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, [])
# User1 now gets no search results for any of the other users.
s = self.get_success(self.handler.search_users(u1, "user2", 10))
@@ -160,14 +160,18 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
r.add((i["user_id"], i["other_user_id"], i["room_id"]))
return r
- def get_users_who_share_public_rooms(self):
- return self.get_success(
+ def get_users_in_public_rooms(self):
+ r = self.get_success(
self.store._simple_select_list(
- "users_who_share_public_rooms",
+ "users_in_public_rooms",
None,
- ["user_id", "other_user_id", "room_id"],
+ ("user_id", "room_id"),
)
)
+ retval = []
+ for i in r:
+ retval.append((i["user_id"], i["room_id"]))
+ return retval
def get_users_who_share_private_rooms(self):
return self.get_success(
@@ -200,11 +204,12 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
self.get_success(self.store.update_user_directory_stream_pos(None))
self.get_success(self.store.delete_all_from_user_dir())
- 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()
+ # Nothing updated yet
self.assertEqual(shares_private, [])
- self.assertEqual(shares_public, [])
+ self.assertEqual(public_users, [])
# Reset the handled users caches
self.handler.initially_handled_users = set()
@@ -219,12 +224,12 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
self.get_success(d)
- 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()
- # User 1 and User 2 share public rooms
+ # User 1 and User 2 are in the same public room
self.assertEqual(
- self._compress_shared(shares_public), set([(u1, u2, room), (u2, u1, room)])
+ set(public_users), set([(u1, room), (u2, room)])
)
# User 1 and User 3 share private rooms
diff --git a/tests/storage/test_user_directory.py b/tests/storage/test_user_directory.py
index a2a652a235..512d76e7a3 100644
--- a/tests/storage/test_user_directory.py
+++ b/tests/storage/test_user_directory.py
@@ -41,8 +41,8 @@ class UserDirectoryStoreTestCase(unittest.TestCase):
BOBBY: ProfileInfo(None, "bobby"),
},
)
- yield self.store.add_users_who_share_room(
- "!room:id", False, ((ALICE, BOB), (BOB, ALICE))
+ yield self.store.add_users_in_public_rooms(
+ "!room:id", (ALICE, BOB)
)
@defer.inlineCallbacks
diff --git a/tests/utils.py b/tests/utils.py
index 9c8dc9dbce..03b5a05b22 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -331,6 +331,8 @@ def setup_test_homeserver(
cleanup_func(cleanup)
hs.setup()
+ if homeserverToUse.__name__ == "TestHomeServer":
+ hs.setup_master()
else:
hs = homeserverToUse(
name,
|
