summary refs log tree commit diff
path: root/tests/storage
diff options
context:
space:
mode:
Diffstat (limited to 'tests/storage')
-rw-r--r--tests/storage/databases/main/test_deviceinbox.py2
-rw-r--r--tests/storage/databases/main/test_events_worker.py6
-rw-r--r--tests/storage/databases/main/test_lock.py2
-rw-r--r--tests/storage/databases/main/test_room.py2
-rw-r--r--tests/storage/test__base.py2
-rw-r--r--tests/storage/test_account_data.py2
-rw-r--r--tests/storage/test_appservice.py2
-rw-r--r--tests/storage/test_background_update.py8
-rw-r--r--tests/storage/test_cleanup_extrems.py4
-rw-r--r--tests/storage/test_client_ips.py4
-rw-r--r--tests/storage/test_devices.py2
-rw-r--r--tests/storage/test_directory.py2
-rw-r--r--tests/storage/test_e2e_room_keys.py2
-rw-r--r--tests/storage/test_end_to_end_keys.py2
-rw-r--r--tests/storage/test_event_chain.py4
-rw-r--r--tests/storage/test_event_federation.py2
-rw-r--r--tests/storage/test_event_push_actions.py2
-rw-r--r--tests/storage/test_events.py4
-rw-r--r--tests/storage/test_id_generators.py6
-rw-r--r--tests/storage/test_keys.py4
-rw-r--r--tests/storage/test_main.py2
-rw-r--r--tests/storage/test_monthly_active_users.py2
-rw-r--r--tests/storage/test_profile.py2
-rw-r--r--tests/storage/test_purge.py4
-rw-r--r--tests/storage/test_redaction.py2
-rw-r--r--tests/storage/test_registration.py2
-rw-r--r--tests/storage/test_rollback_worker.py6
-rw-r--r--tests/storage/test_room.py4
-rw-r--r--tests/storage/test_room_search.py2
-rw-r--r--tests/storage/test_roommember.py4
-rw-r--r--tests/storage/test_state.py2
-rw-r--r--tests/storage/test_stream.py2
-rw-r--r--tests/storage/test_transactions.py2
-rw-r--r--tests/storage/test_user_directory.py4
34 files changed, 52 insertions, 52 deletions
diff --git a/tests/storage/databases/main/test_deviceinbox.py b/tests/storage/databases/main/test_deviceinbox.py
index 36c933b9e9..50c20c5b92 100644
--- a/tests/storage/databases/main/test_deviceinbox.py
+++ b/tests/storage/databases/main/test_deviceinbox.py
@@ -26,7 +26,7 @@ class DeviceInboxBackgroundUpdateStoreTestCase(HomeserverTestCase):
     ]
 
     def prepare(self, reactor, clock, hs):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         self.user_id = self.register_user("foo", "pass")
 
     def test_background_remove_deleted_devices_from_device_inbox(self):
diff --git a/tests/storage/databases/main/test_events_worker.py b/tests/storage/databases/main/test_events_worker.py
index 5ae491ff5a..59def6e59c 100644
--- a/tests/storage/databases/main/test_events_worker.py
+++ b/tests/storage/databases/main/test_events_worker.py
@@ -37,7 +37,7 @@ from tests import unittest
 
 class HaveSeenEventsTestCase(unittest.HomeserverTestCase):
     def prepare(self, reactor, clock, hs):
-        self.store: EventsWorkerStore = hs.get_datastore()
+        self.store: EventsWorkerStore = hs.get_datastores().main
 
         # insert some test data
         for rid in ("room1", "room2"):
@@ -122,7 +122,7 @@ class EventCacheTestCase(unittest.HomeserverTestCase):
     ]
 
     def prepare(self, reactor, clock, hs):
-        self.store: EventsWorkerStore = hs.get_datastore()
+        self.store: EventsWorkerStore = hs.get_datastores().main
 
         self.user = self.register_user("user", "pass")
         self.token = self.login(self.user, "pass")
@@ -163,7 +163,7 @@ class DatabaseOutageTestCase(unittest.HomeserverTestCase):
     """Test event fetching during a database outage."""
 
     def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer):
-        self.store: EventsWorkerStore = hs.get_datastore()
+        self.store: EventsWorkerStore = hs.get_datastores().main
 
         self.room_id = f"!room:{hs.hostname}"
         self.event_ids = [f"event{i}" for i in range(20)]
diff --git a/tests/storage/databases/main/test_lock.py b/tests/storage/databases/main/test_lock.py
index d326a1d6a6..3ac4646969 100644
--- a/tests/storage/databases/main/test_lock.py
+++ b/tests/storage/databases/main/test_lock.py
@@ -20,7 +20,7 @@ from tests import unittest
 
 class LockTestCase(unittest.HomeserverTestCase):
     def prepare(self, reactor, clock, hs: HomeServer):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
 
     def test_simple_lock(self):
         """Test that we can take out a lock and that while we hold it nobody
diff --git a/tests/storage/databases/main/test_room.py b/tests/storage/databases/main/test_room.py
index 7496974da3..9abd0cb446 100644
--- a/tests/storage/databases/main/test_room.py
+++ b/tests/storage/databases/main/test_room.py
@@ -28,7 +28,7 @@ class RoomBackgroundUpdateStoreTestCase(HomeserverTestCase):
     ]
 
     def prepare(self, reactor, clock, hs):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         self.user_id = self.register_user("foo", "pass")
         self.token = self.login("foo", "pass")
 
diff --git a/tests/storage/test__base.py b/tests/storage/test__base.py
index 200b9198f9..4899cd5c36 100644
--- a/tests/storage/test__base.py
+++ b/tests/storage/test__base.py
@@ -20,7 +20,7 @@ from tests import unittest
 
 class UpsertManyTests(unittest.HomeserverTestCase):
     def prepare(self, reactor, clock, hs):
-        self.storage = hs.get_datastore()
+        self.storage = hs.get_datastores().main
 
         self.table_name = "table_" + secrets.token_hex(6)
         self.get_success(
diff --git a/tests/storage/test_account_data.py b/tests/storage/test_account_data.py
index d697d2bc1e..272cd35402 100644
--- a/tests/storage/test_account_data.py
+++ b/tests/storage/test_account_data.py
@@ -21,7 +21,7 @@ from tests import unittest
 
 class IgnoredUsersTestCase(unittest.HomeserverTestCase):
     def prepare(self, hs, reactor, clock):
-        self.store = self.hs.get_datastore()
+        self.store = self.hs.get_datastores().main
         self.user = "@user:test"
 
     def _update_ignore_list(
diff --git a/tests/storage/test_appservice.py b/tests/storage/test_appservice.py
index ddcb7f5549..50703ccaee 100644
--- a/tests/storage/test_appservice.py
+++ b/tests/storage/test_appservice.py
@@ -467,7 +467,7 @@ class ApplicationServiceStoreTypeStreamIds(unittest.HomeserverTestCase):
         self, reactor: MemoryReactor, clock: Clock, homeserver: HomeServer
     ) -> None:
         self.service = Mock(id="foo")
-        self.store = self.hs.get_datastore()
+        self.store = self.hs.get_datastores().main
         self.get_success(
             self.store.set_appservice_state(self.service, ApplicationServiceState.UP)
         )
diff --git a/tests/storage/test_background_update.py b/tests/storage/test_background_update.py
index 6156dfac4e..39dcc094bd 100644
--- a/tests/storage/test_background_update.py
+++ b/tests/storage/test_background_update.py
@@ -24,7 +24,7 @@ from tests.test_utils import make_awaitable, simple_async_mock
 
 class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
     def prepare(self, reactor, clock, homeserver):
-        self.updates: BackgroundUpdater = self.hs.get_datastore().db_pool.updates
+        self.updates: BackgroundUpdater = self.hs.get_datastores().main.db_pool.updates
         # the base test class should have run the real bg updates for us
         self.assertTrue(
             self.get_success(self.updates.has_completed_background_updates())
@@ -42,7 +42,7 @@ class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
         # the target runtime for each bg update
         target_background_update_duration_ms = 100
 
-        store = self.hs.get_datastore()
+        store = self.hs.get_datastores().main
         self.get_success(
             store.db_pool.simple_insert(
                 "background_updates",
@@ -102,7 +102,7 @@ class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
 
 class BackgroundUpdateControllerTestCase(unittest.HomeserverTestCase):
     def prepare(self, reactor, clock, homeserver):
-        self.updates: BackgroundUpdater = self.hs.get_datastore().db_pool.updates
+        self.updates: BackgroundUpdater = self.hs.get_datastores().main.db_pool.updates
         # the base test class should have run the real bg updates for us
         self.assertTrue(
             self.get_success(self.updates.has_completed_background_updates())
@@ -138,7 +138,7 @@ class BackgroundUpdateControllerTestCase(unittest.HomeserverTestCase):
         )
 
     def test_controller(self):
-        store = self.hs.get_datastore()
+        store = self.hs.get_datastores().main
         self.get_success(
             store.db_pool.simple_insert(
                 "background_updates",
diff --git a/tests/storage/test_cleanup_extrems.py b/tests/storage/test_cleanup_extrems.py
index a59c28f896..ce89c96912 100644
--- a/tests/storage/test_cleanup_extrems.py
+++ b/tests/storage/test_cleanup_extrems.py
@@ -30,7 +30,7 @@ class CleanupExtremBackgroundUpdateStoreTestCase(HomeserverTestCase):
     """
 
     def prepare(self, reactor, clock, homeserver):
-        self.store = homeserver.get_datastore()
+        self.store = homeserver.get_datastores().main
         self.room_creator = homeserver.get_room_creation_handler()
 
         # Create a test user and room
@@ -242,7 +242,7 @@ class CleanupExtremDummyEventsTestCase(HomeserverTestCase):
         return self.setup_test_homeserver(config=config)
 
     def prepare(self, reactor, clock, homeserver):
-        self.store = homeserver.get_datastore()
+        self.store = homeserver.get_datastores().main
         self.room_creator = homeserver.get_room_creation_handler()
         self.event_creator_handler = homeserver.get_event_creation_handler()
 
diff --git a/tests/storage/test_client_ips.py b/tests/storage/test_client_ips.py
index c8ac67e35b..49ad3c1324 100644
--- a/tests/storage/test_client_ips.py
+++ b/tests/storage/test_client_ips.py
@@ -35,7 +35,7 @@ class ClientIpStoreTestCase(unittest.HomeserverTestCase):
         return hs
 
     def prepare(self, hs, reactor, clock):
-        self.store = self.hs.get_datastore()
+        self.store = self.hs.get_datastores().main
 
     def test_insert_new_client_ip(self):
         self.reactor.advance(12345678)
@@ -666,7 +666,7 @@ class ClientIpAuthTestCase(unittest.HomeserverTestCase):
         return hs
 
     def prepare(self, hs, reactor, clock):
-        self.store = self.hs.get_datastore()
+        self.store = self.hs.get_datastores().main
         self.user_id = self.register_user("bob", "abc123", True)
 
     def test_request_with_xforwarded(self):
diff --git a/tests/storage/test_devices.py b/tests/storage/test_devices.py
index b547bf8d99..21ffc5a909 100644
--- a/tests/storage/test_devices.py
+++ b/tests/storage/test_devices.py
@@ -19,7 +19,7 @@ from tests.unittest import HomeserverTestCase
 
 class DeviceStoreTestCase(HomeserverTestCase):
     def prepare(self, reactor, clock, hs):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
 
     def test_store_new_device(self):
         self.get_success(
diff --git a/tests/storage/test_directory.py b/tests/storage/test_directory.py
index 43628ce44f..7b72a92424 100644
--- a/tests/storage/test_directory.py
+++ b/tests/storage/test_directory.py
@@ -19,7 +19,7 @@ from tests.unittest import HomeserverTestCase
 
 class DirectoryStoreTestCase(HomeserverTestCase):
     def prepare(self, reactor, clock, hs):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
 
         self.room = RoomID.from_string("!abcde:test")
         self.alias = RoomAlias.from_string("#my-room:test")
diff --git a/tests/storage/test_e2e_room_keys.py b/tests/storage/test_e2e_room_keys.py
index 7556171d8a..fb96ab3a2f 100644
--- a/tests/storage/test_e2e_room_keys.py
+++ b/tests/storage/test_e2e_room_keys.py
@@ -28,7 +28,7 @@ room_key: RoomKey = {
 class E2eRoomKeysHandlerTestCase(unittest.HomeserverTestCase):
     def make_homeserver(self, reactor, clock):
         hs = self.setup_test_homeserver("server", federation_http_client=None)
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         return hs
 
     def test_room_keys_version_delete(self):
diff --git a/tests/storage/test_end_to_end_keys.py b/tests/storage/test_end_to_end_keys.py
index 3bf6e337f4..0f04493ad0 100644
--- a/tests/storage/test_end_to_end_keys.py
+++ b/tests/storage/test_end_to_end_keys.py
@@ -17,7 +17,7 @@ from tests.unittest import HomeserverTestCase
 
 class EndToEndKeyStoreTestCase(HomeserverTestCase):
     def prepare(self, reactor, clock, hs):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
 
     def test_key_without_device_name(self):
         now = 1470174257070
diff --git a/tests/storage/test_event_chain.py b/tests/storage/test_event_chain.py
index e3273a93f9..401020fd63 100644
--- a/tests/storage/test_event_chain.py
+++ b/tests/storage/test_event_chain.py
@@ -30,7 +30,7 @@ from tests.unittest import HomeserverTestCase
 
 class EventChainStoreTestCase(HomeserverTestCase):
     def prepare(self, reactor, clock, hs):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         self._next_stream_ordering = 1
 
     def test_simple(self):
@@ -492,7 +492,7 @@ class EventChainBackgroundUpdateTestCase(HomeserverTestCase):
     ]
 
     def prepare(self, reactor, clock, hs):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         self.user_id = self.register_user("foo", "pass")
         self.token = self.login("foo", "pass")
         self.requester = create_requester(self.user_id)
diff --git a/tests/storage/test_event_federation.py b/tests/storage/test_event_federation.py
index 667ca90a4d..645d564d1c 100644
--- a/tests/storage/test_event_federation.py
+++ b/tests/storage/test_event_federation.py
@@ -31,7 +31,7 @@ import tests.utils
 
 class EventFederationWorkerStoreTestCase(tests.unittest.HomeserverTestCase):
     def prepare(self, reactor, clock, hs):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
 
     def test_get_prev_events_for_room(self):
         room_id = "@ROOM:local"
diff --git a/tests/storage/test_event_push_actions.py b/tests/storage/test_event_push_actions.py
index 738f3ad1dc..c9e3b9fa79 100644
--- a/tests/storage/test_event_push_actions.py
+++ b/tests/storage/test_event_push_actions.py
@@ -30,7 +30,7 @@ HIGHLIGHT = [
 
 class EventPushActionsStoreTestCase(HomeserverTestCase):
     def prepare(self, reactor, clock, hs):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         self.persist_events_store = hs.get_datastores().persist_events
 
     def test_get_unread_push_actions_for_user_in_range_for_http(self):
diff --git a/tests/storage/test_events.py b/tests/storage/test_events.py
index a8639d8f82..ef5e25873c 100644
--- a/tests/storage/test_events.py
+++ b/tests/storage/test_events.py
@@ -32,7 +32,7 @@ class ExtremPruneTestCase(HomeserverTestCase):
     def prepare(self, reactor, clock, homeserver):
         self.state = self.hs.get_state_handler()
         self.persistence = self.hs.get_storage().persistence
-        self.store = self.hs.get_datastore()
+        self.store = self.hs.get_datastores().main
 
         self.register_user("user", "pass")
         self.token = self.login("user", "pass")
@@ -341,7 +341,7 @@ class InvalideUsersInRoomCacheTestCase(HomeserverTestCase):
     def prepare(self, reactor, clock, homeserver):
         self.state = self.hs.get_state_handler()
         self.persistence = self.hs.get_storage().persistence
-        self.store = self.hs.get_datastore()
+        self.store = self.hs.get_datastores().main
 
     def test_remote_user_rooms_cache_invalidated(self):
         """Test that if the server leaves a room the `get_rooms_for_user` cache
diff --git a/tests/storage/test_id_generators.py b/tests/storage/test_id_generators.py
index 7486078284..6ac4b93f98 100644
--- a/tests/storage/test_id_generators.py
+++ b/tests/storage/test_id_generators.py
@@ -26,7 +26,7 @@ class MultiWriterIdGeneratorTestCase(HomeserverTestCase):
         skip = "Requires Postgres"
 
     def prepare(self, reactor, clock, hs):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         self.db_pool: DatabasePool = self.store.db_pool
 
         self.get_success(self.db_pool.runInteraction("_setup_db", self._setup_db))
@@ -459,7 +459,7 @@ class BackwardsMultiWriterIdGeneratorTestCase(HomeserverTestCase):
         skip = "Requires Postgres"
 
     def prepare(self, reactor, clock, hs):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         self.db_pool: DatabasePool = self.store.db_pool
 
         self.get_success(self.db_pool.runInteraction("_setup_db", self._setup_db))
@@ -585,7 +585,7 @@ class MultiTableMultiWriterIdGeneratorTestCase(HomeserverTestCase):
         skip = "Requires Postgres"
 
     def prepare(self, reactor, clock, hs):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         self.db_pool: DatabasePool = self.store.db_pool
 
         self.get_success(self.db_pool.runInteraction("_setup_db", self._setup_db))
diff --git a/tests/storage/test_keys.py b/tests/storage/test_keys.py
index a94b5fd721..9059095525 100644
--- a/tests/storage/test_keys.py
+++ b/tests/storage/test_keys.py
@@ -37,7 +37,7 @@ KEY_2 = decode_verify_key_base64(
 
 class KeyStoreTestCase(tests.unittest.HomeserverTestCase):
     def test_get_server_verify_keys(self):
-        store = self.hs.get_datastore()
+        store = self.hs.get_datastores().main
 
         key_id_1 = "ed25519:key1"
         key_id_2 = "ed25519:KEY_ID_2"
@@ -74,7 +74,7 @@ class KeyStoreTestCase(tests.unittest.HomeserverTestCase):
     def test_cache(self):
         """Check that updates correctly invalidate the cache."""
 
-        store = self.hs.get_datastore()
+        store = self.hs.get_datastores().main
 
         key_id_1 = "ed25519:key1"
         key_id_2 = "ed25519:key2"
diff --git a/tests/storage/test_main.py b/tests/storage/test_main.py
index f8d11bac4e..4ca212fd11 100644
--- a/tests/storage/test_main.py
+++ b/tests/storage/test_main.py
@@ -22,7 +22,7 @@ class DataStoreTestCase(unittest.HomeserverTestCase):
     def setUp(self) -> None:
         super(DataStoreTestCase, self).setUp()
 
-        self.store = self.hs.get_datastore()
+        self.store = self.hs.get_datastores().main
 
         self.user = UserID.from_string("@abcde:test")
         self.displayname = "Frank"
diff --git a/tests/storage/test_monthly_active_users.py b/tests/storage/test_monthly_active_users.py
index d6b4cdd788..79648d45db 100644
--- a/tests/storage/test_monthly_active_users.py
+++ b/tests/storage/test_monthly_active_users.py
@@ -45,7 +45,7 @@ class MonthlyActiveUsersTestCase(unittest.HomeserverTestCase):
         return config
 
     def prepare(self, reactor, clock, homeserver):
-        self.store = homeserver.get_datastore()
+        self.store = homeserver.get_datastores().main
         # Advance the clock a bit
         reactor.advance(FORTY_DAYS)
 
diff --git a/tests/storage/test_profile.py b/tests/storage/test_profile.py
index d37736edf8..b6f99af2f1 100644
--- a/tests/storage/test_profile.py
+++ b/tests/storage/test_profile.py
@@ -22,7 +22,7 @@ from tests import unittest
 
 class ProfileStoreTestCase(unittest.HomeserverTestCase):
     def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
 
         self.u_frank = UserID.from_string("@frank:test")
 
diff --git a/tests/storage/test_purge.py b/tests/storage/test_purge.py
index 22a77c3ccc..08cc60237e 100644
--- a/tests/storage/test_purge.py
+++ b/tests/storage/test_purge.py
@@ -30,7 +30,7 @@ class PurgeTests(HomeserverTestCase):
     def prepare(self, reactor, clock, hs):
         self.room_id = self.helper.create_room_as(self.user_id)
 
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         self.storage = self.hs.get_storage()
 
     def test_purge_history(self):
@@ -47,7 +47,7 @@ class PurgeTests(HomeserverTestCase):
         token = self.get_success(
             self.store.get_topological_token_for_event(last["event_id"])
         )
-        token_str = self.get_success(token.to_string(self.hs.get_datastore()))
+        token_str = self.get_success(token.to_string(self.hs.get_datastores().main))
 
         # Purge everything before this topological token
         self.get_success(
diff --git a/tests/storage/test_redaction.py b/tests/storage/test_redaction.py
index 8c95a0a2fb..03e9cc7d4a 100644
--- a/tests/storage/test_redaction.py
+++ b/tests/storage/test_redaction.py
@@ -30,7 +30,7 @@ class RedactionTestCase(unittest.HomeserverTestCase):
         return config
 
     def prepare(self, reactor, clock, hs):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         self.storage = hs.get_storage()
         self.event_builder_factory = hs.get_event_builder_factory()
         self.event_creation_handler = hs.get_event_creation_handler()
diff --git a/tests/storage/test_registration.py b/tests/storage/test_registration.py
index 9748065282..1fa495f778 100644
--- a/tests/storage/test_registration.py
+++ b/tests/storage/test_registration.py
@@ -20,7 +20,7 @@ from tests.unittest import HomeserverTestCase
 
 class RegistrationStoreTestCase(HomeserverTestCase):
     def prepare(self, reactor, clock, hs):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
 
         self.user_id = "@my-user:test"
         self.tokens = ["AbCdEfGhIjKlMnOpQrStUvWxYz", "BcDeFgHiJkLmNoPqRsTuVwXyZa"]
diff --git a/tests/storage/test_rollback_worker.py b/tests/storage/test_rollback_worker.py
index cfc8098af6..0baa54312e 100644
--- a/tests/storage/test_rollback_worker.py
+++ b/tests/storage/test_rollback_worker.py
@@ -56,7 +56,7 @@ class WorkerSchemaTests(HomeserverTestCase):
     def test_rolling_back(self):
         """Test that workers can start if the DB is a newer schema version"""
 
-        db_pool = self.hs.get_datastore().db_pool
+        db_pool = self.hs.get_datastores().main.db_pool
         db_conn = LoggingDatabaseConnection(
             db_pool._db_pool.connect(),
             db_pool.engine,
@@ -72,7 +72,7 @@ class WorkerSchemaTests(HomeserverTestCase):
 
     def test_not_upgraded_old_schema_version(self):
         """Test that workers don't start if the DB has an older schema version"""
-        db_pool = self.hs.get_datastore().db_pool
+        db_pool = self.hs.get_datastores().main.db_pool
         db_conn = LoggingDatabaseConnection(
             db_pool._db_pool.connect(),
             db_pool.engine,
@@ -92,7 +92,7 @@ class WorkerSchemaTests(HomeserverTestCase):
         Test that workers don't start if the DB is on the current schema version,
         but there are still outstanding delta migrations to run.
         """
-        db_pool = self.hs.get_datastore().db_pool
+        db_pool = self.hs.get_datastores().main.db_pool
         db_conn = LoggingDatabaseConnection(
             db_pool._db_pool.connect(),
             db_pool.engine,
diff --git a/tests/storage/test_room.py b/tests/storage/test_room.py
index 31ce7f6252..42bfca2a83 100644
--- a/tests/storage/test_room.py
+++ b/tests/storage/test_room.py
@@ -23,7 +23,7 @@ class RoomStoreTestCase(HomeserverTestCase):
     def prepare(self, reactor, clock, hs):
         # We can't test RoomStore on its own without the DirectoryStore, for
         # management of the 'room_aliases' table
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
 
         self.room = RoomID.from_string("!abcde:test")
         self.alias = RoomAlias.from_string("#a-room-name:test")
@@ -71,7 +71,7 @@ class RoomEventsStoreTestCase(HomeserverTestCase):
     def prepare(self, reactor, clock, hs):
         # Room events need the full datastore, for persist_event() and
         # get_room_state()
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         self.storage = hs.get_storage()
         self.event_factory = hs.get_event_factory()
 
diff --git a/tests/storage/test_room_search.py b/tests/storage/test_room_search.py
index 8971ecccbd..befaa0fcee 100644
--- a/tests/storage/test_room_search.py
+++ b/tests/storage/test_room_search.py
@@ -46,7 +46,7 @@ class NullByteInsertionTest(HomeserverTestCase):
             self.assertIn("event_id", response)
 
         # Check that search works for the message where the null byte was replaced
-        store = self.hs.get_datastore()
+        store = self.hs.get_datastores().main
         result = self.get_success(
             store.search_msgs([room_id], "hi bob", ["content.body"])
         )
diff --git a/tests/storage/test_roommember.py b/tests/storage/test_roommember.py
index 5cfdfe9b85..7028f0dfb0 100644
--- a/tests/storage/test_roommember.py
+++ b/tests/storage/test_roommember.py
@@ -35,7 +35,7 @@ class RoomMemberStoreTestCase(unittest.HomeserverTestCase):
 
         # We can't test the RoomMemberStore on its own without the other event
         # storage logic
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
 
         self.u_alice = self.register_user("alice", "pass")
         self.t_alice = self.login("alice", "pass")
@@ -212,7 +212,7 @@ class RoomMemberStoreTestCase(unittest.HomeserverTestCase):
 
 class CurrentStateMembershipUpdateTestCase(unittest.HomeserverTestCase):
     def prepare(self, reactor, clock, homeserver):
-        self.store = homeserver.get_datastore()
+        self.store = homeserver.get_datastores().main
         self.room_creator = homeserver.get_room_creation_handler()
 
     def test_can_rerun_update(self):
diff --git a/tests/storage/test_state.py b/tests/storage/test_state.py
index 28c767ecfd..f88f1c55fc 100644
--- a/tests/storage/test_state.py
+++ b/tests/storage/test_state.py
@@ -28,7 +28,7 @@ logger = logging.getLogger(__name__)
 
 class StateStoreTestCase(HomeserverTestCase):
     def prepare(self, reactor, clock, hs):
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
         self.storage = hs.get_storage()
         self.state_datastore = self.storage.state.stores.state
         self.event_builder_factory = hs.get_event_builder_factory()
diff --git a/tests/storage/test_stream.py b/tests/storage/test_stream.py
index ce782c7e1d..6a1cf33054 100644
--- a/tests/storage/test_stream.py
+++ b/tests/storage/test_stream.py
@@ -115,7 +115,7 @@ class PaginationTestCase(HomeserverTestCase):
         )
 
         events, next_key = self.get_success(
-            self.hs.get_datastore().paginate_room_events(
+            self.hs.get_datastores().main.paginate_room_events(
                 room_id=self.room_id,
                 from_key=from_token.room_key,
                 to_key=None,
diff --git a/tests/storage/test_transactions.py b/tests/storage/test_transactions.py
index bea9091d30..e05daa285e 100644
--- a/tests/storage/test_transactions.py
+++ b/tests/storage/test_transactions.py
@@ -20,7 +20,7 @@ from tests.unittest import HomeserverTestCase
 
 class TransactionStoreTestCase(HomeserverTestCase):
     def prepare(self, reactor, clock, homeserver):
-        self.store = homeserver.get_datastore()
+        self.store = homeserver.get_datastores().main
 
     def test_get_set_transactions(self):
         """Tests that we can successfully get a non-existent entry for
diff --git a/tests/storage/test_user_directory.py b/tests/storage/test_user_directory.py
index 48f1e9d841..7f1964eb6a 100644
--- a/tests/storage/test_user_directory.py
+++ b/tests/storage/test_user_directory.py
@@ -149,7 +149,7 @@ class UserDirectoryInitialPopulationTestcase(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.user_dir_helper = GetUserDirectoryTables(self.store)
 
     def _purge_and_rebuild_user_dir(self) -> None:
@@ -415,7 +415,7 @@ class UserDirectoryInitialPopulationTestcase(HomeserverTestCase):
 
 class UserDirectoryStoreTestCase(HomeserverTestCase):
     def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
-        self.store = hs.get_datastore()
+        self.store = hs.get_datastores().main
 
         # alice and bob are both in !room_id. bobby is not but shares
         # a homeserver with alice.