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.
|
