diff --git a/tests/replication/_base.py b/tests/replication/_base.py
index 9fc50f8852..a7a05a564f 100644
--- a/tests/replication/_base.py
+++ b/tests/replication/_base.py
@@ -68,7 +68,7 @@ class BaseStreamTestCase(unittest.HomeserverTestCase):
# Since we use sqlite in memory databases we need to make sure the
# databases objects are the same.
- self.worker_hs.get_datastore().db_pool = hs.get_datastore().db_pool
+ self.worker_hs.get_datastores().main.db_pool = hs.get_datastores().main.db_pool
# Normally we'd pass in the handler to `setup_test_homeserver`, which would
# eventually hit "Install @cache_in_self attributes" in tests/utils.py.
@@ -233,7 +233,7 @@ class BaseMultiWorkerStreamTestCase(unittest.HomeserverTestCase):
# We may have an attempt to connect to redis for the external cache already.
self.connect_any_redis_attempts()
- store = self.hs.get_datastore()
+ store = self.hs.get_datastores().main
self.database_pool = store.db_pool
self.reactor.lookups["testserv"] = "1.2.3.4"
@@ -332,7 +332,7 @@ class BaseMultiWorkerStreamTestCase(unittest.HomeserverTestCase):
lambda: self._handle_http_replication_attempt(worker_hs, port),
)
- store = worker_hs.get_datastore()
+ store = worker_hs.get_datastores().main
store.db_pool._db_pool = self.database_pool._db_pool
# Set up TCP replication between master and the new worker if we don't
diff --git a/tests/replication/slave/storage/_base.py b/tests/replication/slave/storage/_base.py
index 83e89383f6..85be79d19d 100644
--- a/tests/replication/slave/storage/_base.py
+++ b/tests/replication/slave/storage/_base.py
@@ -30,8 +30,8 @@ class BaseSlavedStoreTestCase(BaseStreamTestCase):
self.reconnect()
- self.master_store = hs.get_datastore()
- self.slaved_store = self.worker_hs.get_datastore()
+ self.master_store = hs.get_datastores().main
+ self.slaved_store = self.worker_hs.get_datastores().main
self.storage = hs.get_storage()
def replicate(self):
diff --git a/tests/replication/tcp/streams/test_account_data.py b/tests/replication/tcp/streams/test_account_data.py
index cdd052001b..50fbff5f32 100644
--- a/tests/replication/tcp/streams/test_account_data.py
+++ b/tests/replication/tcp/streams/test_account_data.py
@@ -23,7 +23,7 @@ from tests.replication._base import BaseStreamTestCase
class AccountDataStreamTestCase(BaseStreamTestCase):
def test_update_function_room_account_data_limit(self):
"""Test replication with many room account data updates"""
- store = self.hs.get_datastore()
+ store = self.hs.get_datastores().main
# generate lots of account data updates
updates = []
@@ -69,7 +69,7 @@ class AccountDataStreamTestCase(BaseStreamTestCase):
def test_update_function_global_account_data_limit(self):
"""Test replication with many global account data updates"""
- store = self.hs.get_datastore()
+ store = self.hs.get_datastores().main
# generate lots of account data updates
updates = []
diff --git a/tests/replication/tcp/streams/test_events.py b/tests/replication/tcp/streams/test_events.py
index f198a94887..f9d5da723c 100644
--- a/tests/replication/tcp/streams/test_events.py
+++ b/tests/replication/tcp/streams/test_events.py
@@ -136,7 +136,7 @@ class EventsStreamTestCase(BaseStreamTestCase):
# this is the point in the DAG where we make a fork
fork_point: List[str] = self.get_success(
- self.hs.get_datastore().get_latest_event_ids_in_room(self.room_id)
+ self.hs.get_datastores().main.get_latest_event_ids_in_room(self.room_id)
)
events = [
@@ -291,7 +291,7 @@ class EventsStreamTestCase(BaseStreamTestCase):
# this is the point in the DAG where we make a fork
fork_point: List[str] = self.get_success(
- self.hs.get_datastore().get_latest_event_ids_in_room(self.room_id)
+ self.hs.get_datastores().main.get_latest_event_ids_in_room(self.room_id)
)
events: List[EventBase] = []
diff --git a/tests/replication/tcp/streams/test_receipts.py b/tests/replication/tcp/streams/test_receipts.py
index 38e292c1ab..eb00117845 100644
--- a/tests/replication/tcp/streams/test_receipts.py
+++ b/tests/replication/tcp/streams/test_receipts.py
@@ -32,7 +32,7 @@ class ReceiptsStreamTestCase(BaseStreamTestCase):
# tell the master to send a new receipt
self.get_success(
- self.hs.get_datastore().insert_receipt(
+ self.hs.get_datastores().main.insert_receipt(
"!room:blue", "m.read", USER_ID, ["$event:blue"], {"a": 1}
)
)
@@ -56,7 +56,7 @@ class ReceiptsStreamTestCase(BaseStreamTestCase):
self.test_handler.on_rdata.reset_mock()
self.get_success(
- self.hs.get_datastore().insert_receipt(
+ self.hs.get_datastores().main.insert_receipt(
"!room2:blue", "m.read", USER_ID, ["$event2:foo"], {"a": 2}
)
)
diff --git a/tests/replication/test_federation_sender_shard.py b/tests/replication/test_federation_sender_shard.py
index 92a5b53e11..ba1a63c0d6 100644
--- a/tests/replication/test_federation_sender_shard.py
+++ b/tests/replication/test_federation_sender_shard.py
@@ -204,7 +204,7 @@ class FederationSenderTestCase(BaseMultiWorkerStreamTestCase):
def create_room_with_remote_server(self, user, token, remote_server="other_server"):
room = self.helper.create_room_as(user, tok=token)
- store = self.hs.get_datastore()
+ store = self.hs.get_datastores().main
federation = self.hs.get_federation_event_handler()
prev_event_ids = self.get_success(store.get_latest_event_ids_in_room(room))
diff --git a/tests/replication/test_pusher_shard.py b/tests/replication/test_pusher_shard.py
index 4094a75f36..8f4f6688ce 100644
--- a/tests/replication/test_pusher_shard.py
+++ b/tests/replication/test_pusher_shard.py
@@ -50,7 +50,7 @@ class PusherShardTestCase(BaseMultiWorkerStreamTestCase):
# Register a pusher
user_dict = self.get_success(
- self.hs.get_datastore().get_user_by_access_token(access_token)
+ self.hs.get_datastores().main.get_user_by_access_token(access_token)
)
token_id = user_dict.token_id
diff --git a/tests/replication/test_sharded_event_persister.py b/tests/replication/test_sharded_event_persister.py
index 596ba5a0c9..5f142e84c3 100644
--- a/tests/replication/test_sharded_event_persister.py
+++ b/tests/replication/test_sharded_event_persister.py
@@ -47,7 +47,7 @@ class EventPersisterShardTestCase(BaseMultiWorkerStreamTestCase):
self.other_access_token = self.login("otheruser", "pass")
self.room_creator = self.hs.get_room_creation_handler()
- self.store = hs.get_datastore()
+ self.store = hs.get_datastores().main
def default_config(self):
conf = super().default_config()
@@ -99,7 +99,7 @@ class EventPersisterShardTestCase(BaseMultiWorkerStreamTestCase):
persisted_on_1 = False
persisted_on_2 = False
- store = self.hs.get_datastore()
+ store = self.hs.get_datastores().main
user_id = self.register_user("user", "pass")
access_token = self.login("user", "pass")
@@ -166,7 +166,7 @@ class EventPersisterShardTestCase(BaseMultiWorkerStreamTestCase):
user_id = self.register_user("user", "pass")
access_token = self.login("user", "pass")
- store = self.hs.get_datastore()
+ store = self.hs.get_datastores().main
# Create two room on the different workers.
self._create_room(room_id1, user_id, access_token)
@@ -194,7 +194,7 @@ class EventPersisterShardTestCase(BaseMultiWorkerStreamTestCase):
#
# Worker2's event stream position will not advance until we call
# __aexit__ again.
- worker_store2 = worker_hs2.get_datastore()
+ worker_store2 = worker_hs2.get_datastores().main
assert isinstance(worker_store2._stream_id_gen, MultiWriterIdGenerator)
actx = worker_store2._stream_id_gen.get_next()
|
