Remove `HomeServer.get_datastore()` (#12031)

The presence of this method was confusing, and mostly present for backwards
compatibility. Let's get rid of it.

Part of #11733

18 files changed, 39 insertions, 37 deletions

diff --git a/tests/handlers/test_appservice.py b/tests/handlers/test_appservice.py
index fe57ff2671..9918ff6807 100644
--- a/tests/handlers/test_appservice.py
+++ b/tests/handlers/test_appservice.py
@@ -38,7 +38,7 @@ class AppServiceHandlerTestCase(unittest.TestCase):
         self.mock_as_api = Mock()
         self.mock_scheduler = Mock()
         hs = Mock()
-        hs.get_datastore.return_value = self.mock_store
+        hs.get_datastores.return_value = Mock(main=self.mock_store)
         self.mock_store.get_received_ts.return_value = make_awaitable(0)
         self.mock_store.set_appservice_last_pos.return_value = make_awaitable(None)
         self.mock_store.set_appservice_stream_type_pos.return_value = make_awaitable(
@@ -355,7 +355,9 @@ class ApplicationServicesHandlerSendEventsTestCase(unittest.HomeserverTestCase):
 
         # Mock out application services, and allow defining our own in tests
         self._services: List[ApplicationService] = []
-        self.hs.get_datastore().get_app_services = Mock(return_value=self._services)
+        self.hs.get_datastores().main.get_app_services = Mock(
+            return_value=self._services
+        )
 
         # A user on the homeserver.
         self.local_user_device_id = "local_device"
@@ -494,7 +496,7 @@ class ApplicationServicesHandlerSendEventsTestCase(unittest.HomeserverTestCase):
         # Create a fake device per message. We can't send to-device messages to
         # a device that doesn't exist.
         self.get_success(
-            self.hs.get_datastore().db_pool.simple_insert_many(
+            self.hs.get_datastores().main.db_pool.simple_insert_many(
                 desc="test_application_services_receive_burst_of_to_device",
                 table="devices",
                 keys=("user_id", "device_id"),
@@ -510,7 +512,7 @@ class ApplicationServicesHandlerSendEventsTestCase(unittest.HomeserverTestCase):
 
         # Seed the device_inbox table with our fake messages
         self.get_success(
-            self.hs.get_datastore().add_messages_to_device_inbox(messages, {})
+            self.hs.get_datastores().main.add_messages_to_device_inbox(messages, {})
         )
 
         # Now have local_user send a final to-device message to exclusive_as_user. All unsent
diff --git a/tests/handlers/test_auth.py b/tests/handlers/test_auth.py
index 03b8b8615c..0c6e55e725 100644
--- a/tests/handlers/test_auth.py
+++ b/tests/handlers/test_auth.py
@@ -129,7 +129,7 @@ class AuthTestCase(unittest.HomeserverTestCase):
 
     def test_mau_limits_exceeded_large(self):
         self.auth_blocking._limit_usage_by_mau = True
-        self.hs.get_datastore().get_monthly_active_count = Mock(
+        self.hs.get_datastores().main.get_monthly_active_count = Mock(
             return_value=make_awaitable(self.large_number_of_users)
         )
 
@@ -140,7 +140,7 @@ class AuthTestCase(unittest.HomeserverTestCase):
             ResourceLimitError,
         )
 
-        self.hs.get_datastore().get_monthly_active_count = Mock(
+        self.hs.get_datastores().main.get_monthly_active_count = Mock(
             return_value=make_awaitable(self.large_number_of_users)
         )
         self.get_failure(
@@ -156,7 +156,7 @@ class AuthTestCase(unittest.HomeserverTestCase):
         self.auth_blocking._limit_usage_by_mau = True
 
         # Set the server to be at the edge of too many users.
-        self.hs.get_datastore().get_monthly_active_count = Mock(
+        self.hs.get_datastores().main.get_monthly_active_count = Mock(
             return_value=make_awaitable(self.auth_blocking._max_mau_value)
         )
 
@@ -175,7 +175,7 @@ class AuthTestCase(unittest.HomeserverTestCase):
         )
 
         # If in monthly active cohort
-        self.hs.get_datastore().user_last_seen_monthly_active = Mock(
+        self.hs.get_datastores().main.user_last_seen_monthly_active = Mock(
             return_value=make_awaitable(self.clock.time_msec())
         )
         self.get_success(
@@ -192,7 +192,7 @@ class AuthTestCase(unittest.HomeserverTestCase):
     def test_mau_limits_not_exceeded(self):
         self.auth_blocking._limit_usage_by_mau = True
 
-        self.hs.get_datastore().get_monthly_active_count = Mock(
+        self.hs.get_datastores().main.get_monthly_active_count = Mock(
             return_value=make_awaitable(self.small_number_of_users)
         )
         # Ensure does not raise exception
@@ -202,7 +202,7 @@ class AuthTestCase(unittest.HomeserverTestCase):
             )
         )
 
-        self.hs.get_datastore().get_monthly_active_count = Mock(
+        self.hs.get_datastores().main.get_monthly_active_count = Mock(
             return_value=make_awaitable(self.small_number_of_users)
         )
         self.get_success(
diff --git a/tests/handlers/test_cas.py b/tests/handlers/test_cas.py
index 8705ff8943..a267228846 100644
--- a/tests/handlers/test_cas.py
+++ b/tests/handlers/test_cas.py
@@ -77,7 +77,7 @@ class CasHandlerTestCase(HomeserverTestCase):
 
     def test_map_cas_user_to_existing_user(self):
         """Existing users can log in with CAS account."""
-        store = self.hs.get_datastore()
+        store = self.hs.get_datastores().main
         self.get_success(
             store.register_user(user_id="@test_user:test", password_hash=None)
         )
diff --git a/tests/handlers/test_deactivate_account.py b/tests/handlers/test_deactivate_account.py
index 01096a1581..ddda36c5a9 100644
--- a/tests/handlers/test_deactivate_account.py
+++ b/tests/handlers/test_deactivate_account.py
@@ -34,7 +34,7 @@ class DeactivateAccountTestCase(HomeserverTestCase):
     ]
 
     def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
-        self._store = hs.get_datastore()
+        self._store = hs.get_datastores().main
 
         self.user = self.register_user("user", "pass")
         self.token = self.login("user", "pass")
diff --git a/tests/handlers/test_device.py b/tests/handlers/test_device.py
index 43031e07ea..683677fd07 100644
--- a/tests/handlers/test_device.py
+++ b/tests/handlers/test_device.py
@@ -28,7 +28,7 @@ class DeviceTestCase(unittest.HomeserverTestCase):
     def make_homeserver(self, reactor, clock):
         hs = self.setup_test_homeserver("server", federation_http_client=None)
         self.handler = hs.get_device_handler()
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         return hs
 
     def prepare(self, reactor, clock, hs):
@@ -263,7 +263,7 @@ class DehydrationTestCase(unittest.HomeserverTestCase):
         self.handler = hs.get_device_handler()
         self.registration = hs.get_registration_handler()
         self.auth = hs.get_auth()
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         return hs
 
     def test_dehydrate_and_rehydrate_device(self):
diff --git a/tests/handlers/test_directory.py b/tests/handlers/test_directory.py
index 0ea4e753e2..65ab107d0e 100644
--- a/tests/handlers/test_directory.py
+++ b/tests/handlers/test_directory.py
@@ -46,7 +46,7 @@ class DirectoryTestCase(unittest.HomeserverTestCase):
 
         self.handler = hs.get_directory_handler()
 
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
 
         self.my_room = RoomAlias.from_string("#my-room:test")
         self.your_room = RoomAlias.from_string("#your-room:test")
@@ -174,7 +174,7 @@ class TestDeleteAlias(unittest.HomeserverTestCase):
     ]
 
     def prepare(self, reactor, clock, hs):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         self.handler = hs.get_directory_handler()
         self.state_handler = hs.get_state_handler()
 
@@ -289,7 +289,7 @@ class CanonicalAliasTestCase(unittest.HomeserverTestCase):
     ]
 
     def prepare(self, reactor, clock, hs):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         self.handler = hs.get_directory_handler()
         self.state_handler = hs.get_state_handler()
 
diff --git a/tests/handlers/test_e2e_keys.py b/tests/handlers/test_e2e_keys.py
index 734ed84d78..9338ab92e9 100644
--- a/tests/handlers/test_e2e_keys.py
+++ b/tests/handlers/test_e2e_keys.py
@@ -34,7 +34,7 @@ class E2eKeysHandlerTestCase(unittest.HomeserverTestCase):
 
     def prepare(self, reactor, clock, hs):
         self.handler = hs.get_e2e_keys_handler()
-        self.store = self.hs.get_datastore()
+        self.store = self.hs.get_datastores().main
 
     def test_query_local_devices_no_devices(self):
         """If the user has no devices, we expect an empty list."""
diff --git a/tests/handlers/test_federation.py b/tests/handlers/test_federation.py
index 496b581726..e8b4e39d1a 100644
--- a/tests/handlers/test_federation.py
+++ b/tests/handlers/test_federation.py
@@ -45,7 +45,7 @@ class FederationTestCase(unittest.HomeserverTestCase):
     def make_homeserver(self, reactor, clock):
         hs = self.setup_test_homeserver(federation_http_client=None)
         self.handler = hs.get_federation_handler()
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         self.state_store = hs.get_storage().state
         self._event_auth_handler = hs.get_event_auth_handler()
         return hs
diff --git a/tests/handlers/test_message.py b/tests/handlers/test_message.py
index 5816295d8b..f4f7ab4845 100644
--- a/tests/handlers/test_message.py
+++ b/tests/handlers/test_message.py
@@ -44,7 +44,7 @@ class EventCreationTestCase(unittest.HomeserverTestCase):
         self.room_id = self.helper.create_room_as(self.user_id, tok=self.access_token)
 
         self.info = self.get_success(
-            self.hs.get_datastore().get_user_by_access_token(
+            self.hs.get_datastores().main.get_user_by_access_token(
                 self.access_token,
             )
         )
diff --git a/tests/handlers/test_oidc.py b/tests/handlers/test_oidc.py
index a552d8182e..e8418b6638 100644
--- a/tests/handlers/test_oidc.py
+++ b/tests/handlers/test_oidc.py
@@ -856,7 +856,7 @@ class OidcHandlerTestCase(HomeserverTestCase):
         auth_handler.complete_sso_login.reset_mock()
 
         # Test if the mxid is already taken
-        store = self.hs.get_datastore()
+        store = self.hs.get_datastores().main
         user3 = UserID.from_string("@test_user_3:test")
         self.get_success(
             store.register_user(user_id=user3.to_string(), password_hash=None)
@@ -872,7 +872,7 @@ class OidcHandlerTestCase(HomeserverTestCase):
     @override_config({"oidc_config": {**DEFAULT_CONFIG, "allow_existing_users": True}})
     def test_map_userinfo_to_existing_user(self):
         """Existing users can log in with OpenID Connect when allow_existing_users is True."""
-        store = self.hs.get_datastore()
+        store = self.hs.get_datastores().main
         user = UserID.from_string("@test_user:test")
         self.get_success(
             store.register_user(user_id=user.to_string(), password_hash=None)
@@ -996,7 +996,7 @@ class OidcHandlerTestCase(HomeserverTestCase):
         auth_handler = self.hs.get_auth_handler()
         auth_handler.complete_sso_login = simple_async_mock()
 
-        store = self.hs.get_datastore()
+        store = self.hs.get_datastores().main
         self.get_success(
             store.register_user(user_id="@test_user:test", password_hash=None)
         )
diff --git a/tests/handlers/test_presence.py b/tests/handlers/test_presence.py
index 671dc7d083..61d28603ae 100644
--- a/tests/handlers/test_presence.py
+++ b/tests/handlers/test_presence.py
@@ -43,7 +43,7 @@ class PresenceUpdateTestCase(unittest.HomeserverTestCase):
     servlets = [admin.register_servlets]
 
     def prepare(self, reactor, clock, homeserver):
-        self.store = homeserver.get_datastore()
+        self.store = homeserver.get_datastores().main
 
     def test_offline_to_online(self):
         wheel_timer = Mock()
@@ -891,7 +891,7 @@ class PresenceJoinTestCase(unittest.HomeserverTestCase):
         # self.event_builder_for_2 = EventBuilderFactory(hs)
         # self.event_builder_for_2.hostname = "test2"
 
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         self.state = hs.get_state_handler()
         self._event_auth_handler = hs.get_event_auth_handler()
 
diff --git a/tests/handlers/test_profile.py b/tests/handlers/test_profile.py
index 60235e5699..69e299fc17 100644
--- a/tests/handlers/test_profile.py
+++ b/tests/handlers/test_profile.py
@@ -48,7 +48,7 @@ class ProfileTestCase(unittest.HomeserverTestCase):
         return hs
 
     def prepare(self, reactor, clock, hs: HomeServer):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
 
         self.frank = UserID.from_string("@1234abcd:test")
         self.bob = UserID.from_string("@4567:test")
@@ -325,7 +325,7 @@ class ProfileTestCase(unittest.HomeserverTestCase):
                 properties are "mimetype" (for the file's type) and "size" (for the
                 file's size).
         """
-        store = self.hs.get_datastore()
+        store = self.hs.get_datastores().main
 
         for name, props in names_and_props.items():
             self.get_success(
diff --git a/tests/handlers/test_register.py b/tests/handlers/test_register.py
index cd6f2c77ae..51ee667ab4 100644
--- a/tests/handlers/test_register.py
+++ b/tests/handlers/test_register.py
@@ -154,7 +154,7 @@ class RegistrationTestCase(unittest.HomeserverTestCase):
 
     def prepare(self, reactor, clock, hs):
         self.handler = self.hs.get_registration_handler()
-        self.store = self.hs.get_datastore()
+        self.store = self.hs.get_datastores().main
         self.lots_of_users = 100
         self.small_number_of_users = 1
 
@@ -172,7 +172,7 @@ class RegistrationTestCase(unittest.HomeserverTestCase):
         self.assertGreater(len(result_token), 20)
 
     def test_if_user_exists(self):
-        store = self.hs.get_datastore()
+        store = self.hs.get_datastores().main
         frank = UserID.from_string("@frank:test")
         self.get_success(
             store.register_user(user_id=frank.to_string(), password_hash=None)
@@ -760,7 +760,7 @@ class RemoteAutoJoinTestCase(unittest.HomeserverTestCase):
 
     def prepare(self, reactor, clock, hs):
         self.handler = self.hs.get_registration_handler()
-        self.store = self.hs.get_datastore()
+        self.store = self.hs.get_datastores().main
 
     @override_config({"auto_join_rooms": ["#room:remotetest"]})
     def test_auto_create_auto_join_remote_room(self):
diff --git a/tests/handlers/test_saml.py b/tests/handlers/test_saml.py
index 50551aa6e3..23941abed8 100644
--- a/tests/handlers/test_saml.py
+++ b/tests/handlers/test_saml.py
@@ -142,7 +142,7 @@ class SamlHandlerTestCase(HomeserverTestCase):
     @override_config({"saml2_config": {"grandfathered_mxid_source_attribute": "mxid"}})
     def test_map_saml_response_to_existing_user(self):
         """Existing users can log in with SAML account."""
-        store = self.hs.get_datastore()
+        store = self.hs.get_datastores().main
         self.get_success(
             store.register_user(user_id="@test_user:test", password_hash=None)
         )
@@ -217,7 +217,7 @@ class SamlHandlerTestCase(HomeserverTestCase):
         sso_handler.render_error = Mock(return_value=None)
 
         # register a user to occupy the first-choice MXID
-        store = self.hs.get_datastore()
+        store = self.hs.get_datastores().main
         self.get_success(
             store.register_user(user_id="@test_user:test", password_hash=None)
         )
diff --git a/tests/handlers/test_stats.py b/tests/handlers/test_stats.py
index 56207f4db6..ecd78fa369 100644
--- a/tests/handlers/test_stats.py
+++ b/tests/handlers/test_stats.py
@@ -33,7 +33,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
     ]
 
     def prepare(self, reactor, clock, hs):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         self.handler = self.hs.get_stats_handler()
 
     def _add_background_updates(self):
diff --git a/tests/handlers/test_sync.py b/tests/handlers/test_sync.py
index 07a760e91a..66b0bd4d1a 100644
--- a/tests/handlers/test_sync.py
+++ b/tests/handlers/test_sync.py
@@ -41,7 +41,7 @@ class SyncTestCase(tests.unittest.HomeserverTestCase):
 
     def prepare(self, reactor, clock, hs: HomeServer):
         self.sync_handler = self.hs.get_sync_handler()
-        self.store = self.hs.get_datastore()
+        self.store = self.hs.get_datastores().main
 
         # AuthBlocking reads from the hs' config on initialization. We need to
         # modify its config instead of the hs'
@@ -248,7 +248,7 @@ class SyncTestCase(tests.unittest.HomeserverTestCase):
         # the prev_events used when creating the join event, such that the ban does not
         # precede the join.
         mocked_get_prev_events = patch.object(
-            self.hs.get_datastore(),
+            self.hs.get_datastores().main,
             "get_prev_events_for_room",
             new_callable=MagicMock,
             return_value=make_awaitable([last_room_creation_event_id]),
diff --git a/tests/handlers/test_typing.py b/tests/handlers/test_typing.py
index 000f9b9fde..e461e03599 100644
--- a/tests/handlers/test_typing.py
+++ b/tests/handlers/test_typing.py
@@ -91,7 +91,7 @@ class TypingNotificationsTestCase(unittest.HomeserverTestCase):
 
         self.event_source = hs.get_event_sources().sources.typing
 
-        self.datastore = hs.get_datastore()
+        self.datastore = hs.get_datastores().main
         self.datastore.get_destination_retry_timings = Mock(
             return_value=defer.succeed(None)
         )
diff --git a/tests/handlers/test_user_directory.py b/tests/handlers/test_user_directory.py
index 482c90ef68..e159169e22 100644
--- a/tests/handlers/test_user_directory.py
+++ b/tests/handlers/test_user_directory.py
@@ -77,7 +77,7 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
         return hs
 
     def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         self.handler = hs.get_user_directory_handler()
         self.event_builder_factory = self.hs.get_event_builder_factory()
         self.event_creation_handler = self.hs.get_event_creation_handler()