diff options
Diffstat (limited to 'tests/storage/test_user_directory.py')
| -rw-r--r-- | tests/storage/test_user_directory.py | 192 |
1 files changed, 188 insertions, 4 deletions
diff --git a/tests/storage/test_user_directory.py b/tests/storage/test_user_directory.py index 222e5d129d..74c8a8599e 100644 --- a/tests/storage/test_user_directory.py +++ b/tests/storage/test_user_directory.py @@ -11,6 +11,15 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from typing import Dict, List, Set, Tuple + +from twisted.test.proto_helpers import MemoryReactor + +from synapse.rest import admin +from synapse.rest.client import login, room +from synapse.server import HomeServer +from synapse.storage import DataStore +from synapse.util import Clock from tests.unittest import HomeserverTestCase, override_config @@ -21,8 +30,183 @@ BOBBY = "@bobby:a" BELA = "@somenickname:a" +class GetUserDirectoryTables: + """Helper functions that we want to reuse in tests/handlers/test_user_directory.py""" + + def __init__(self, store: DataStore): + self.store = store + + def _compress_shared( + self, shared: List[Dict[str, str]] + ) -> Set[Tuple[str, str, str]]: + """ + Compress a list of users who share rooms dicts to a list of tuples. + """ + r = set() + for i in shared: + r.add((i["user_id"], i["other_user_id"], i["room_id"])) + return r + + async def get_users_in_public_rooms(self) -> List[Tuple[str, str]]: + r = await self.store.db_pool.simple_select_list( + "users_in_public_rooms", None, ("user_id", "room_id") + ) + + retval = [] + for i in r: + retval.append((i["user_id"], i["room_id"])) + return retval + + async def get_users_who_share_private_rooms(self) -> List[Dict[str, str]]: + return await self.store.db_pool.simple_select_list( + "users_who_share_private_rooms", + None, + ["user_id", "other_user_id", "room_id"], + ) + + +class UserDirectoryInitialPopulationTestcase(HomeserverTestCase): + """Ensure that rebuilding the directory writes the correct data to the DB. + + See also tests/handlers/test_user_directory.py for similar checks. They + test the incremental updates, rather than the big rebuild. + """ + + servlets = [ + login.register_servlets, + admin.register_servlets_for_client_rest_resource, + room.register_servlets, + ] + + def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None: + self.store = hs.get_datastore() + self.user_dir_helper = GetUserDirectoryTables(self.store) + + def _purge_and_rebuild_user_dir(self) -> None: + """Nuke the user directory tables, start the background process to + repopulate them, and wait for the process to complete. This allows us + to inspect the outcome of the background process alone, without any of + the other incremental updates. + """ + self.get_success(self.store.update_user_directory_stream_pos(None)) + self.get_success(self.store.delete_all_from_user_dir()) + + shares_private = self.get_success( + self.user_dir_helper.get_users_who_share_private_rooms() + ) + public_users = self.get_success( + self.user_dir_helper.get_users_in_public_rooms() + ) + + # Nothing updated yet + self.assertEqual(shares_private, []) + self.assertEqual(public_users, []) + + # Ugh, have to reset this flag + self.store.db_pool.updates._all_done = False + + self.get_success( + self.store.db_pool.simple_insert( + "background_updates", + { + "update_name": "populate_user_directory_createtables", + "progress_json": "{}", + }, + ) + ) + self.get_success( + self.store.db_pool.simple_insert( + "background_updates", + { + "update_name": "populate_user_directory_process_rooms", + "progress_json": "{}", + "depends_on": "populate_user_directory_createtables", + }, + ) + ) + self.get_success( + self.store.db_pool.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.db_pool.simple_insert( + "background_updates", + { + "update_name": "populate_user_directory_cleanup", + "progress_json": "{}", + "depends_on": "populate_user_directory_process_users", + }, + ) + ) + + while not self.get_success( + self.store.db_pool.updates.has_completed_background_updates() + ): + self.get_success( + self.store.db_pool.updates.do_next_background_update(100), by=0.1 + ) + + def test_initial(self) -> None: + """ + The user directory's initial handler correctly updates the search tables. + """ + u1 = self.register_user("user1", "pass") + u1_token = self.login(u1, "pass") + u2 = self.register_user("user2", "pass") + u2_token = self.login(u2, "pass") + u3 = self.register_user("user3", "pass") + u3_token = self.login(u3, "pass") + + 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) + + private_room = self.helper.create_room_as(u1, is_public=False, tok=u1_token) + self.helper.invite(private_room, src=u1, targ=u3, tok=u1_token) + self.helper.join(private_room, user=u3, tok=u3_token) + + self.get_success(self.store.update_user_directory_stream_pos(None)) + self.get_success(self.store.delete_all_from_user_dir()) + + shares_private = self.get_success( + self.user_dir_helper.get_users_who_share_private_rooms() + ) + public_users = self.get_success( + self.user_dir_helper.get_users_in_public_rooms() + ) + + # Nothing updated yet + self.assertEqual(shares_private, []) + self.assertEqual(public_users, []) + + # Do the initial population of the user directory via the background update + self._purge_and_rebuild_user_dir() + + shares_private = self.get_success( + self.user_dir_helper.get_users_who_share_private_rooms() + ) + public_users = self.get_success( + self.user_dir_helper.get_users_in_public_rooms() + ) + + # User 1 and User 2 are in the same public room + self.assertEqual(set(public_users), {(u1, room), (u2, room)}) + + # User 1 and User 3 share private rooms + self.assertEqual( + self.user_dir_helper._compress_shared(shares_private), + {(u1, u3, private_room), (u3, u1, private_room)}, + ) + + class UserDirectoryStoreTestCase(HomeserverTestCase): - def prepare(self, reactor, clock, hs): + def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None: self.store = hs.get_datastore() # alice and bob are both in !room_id. bobby is not but shares @@ -33,7 +217,7 @@ class UserDirectoryStoreTestCase(HomeserverTestCase): self.get_success(self.store.update_profile_in_user_dir(BELA, "Bela", None)) self.get_success(self.store.add_users_in_public_rooms("!room:id", (ALICE, BOB))) - def test_search_user_dir(self): + def test_search_user_dir(self) -> None: # normally when alice searches the directory she should just find # bob because bobby doesn't share a room with her. r = self.get_success(self.store.search_user_dir(ALICE, "bob", 10)) @@ -44,7 +228,7 @@ class UserDirectoryStoreTestCase(HomeserverTestCase): ) @override_config({"user_directory": {"search_all_users": True}}) - def test_search_user_dir_all_users(self): + def test_search_user_dir_all_users(self) -> None: r = self.get_success(self.store.search_user_dir(ALICE, "bob", 10)) self.assertFalse(r["limited"]) self.assertEqual(2, len(r["results"])) @@ -58,7 +242,7 @@ class UserDirectoryStoreTestCase(HomeserverTestCase): ) @override_config({"user_directory": {"search_all_users": True}}) - def test_search_user_dir_stop_words(self): + def test_search_user_dir_stop_words(self) -> None: """Tests that a user can look up another user by searching for the start if its display name even if that name happens to be a common English word that would usually be ignored in full text searches. |