summary refs log tree commit diff
path: root/tests/replication
diff options
context:
space:
mode:
Diffstat (limited to 'tests/replication')
-rw-r--r--tests/replication/_base.py6
-rw-r--r--tests/replication/slave/storage/_base.py4
-rw-r--r--tests/replication/tcp/streams/test_account_data.py4
-rw-r--r--tests/replication/tcp/streams/test_events.py4
-rw-r--r--tests/replication/tcp/streams/test_receipts.py4
-rw-r--r--tests/replication/test_federation_sender_shard.py2
-rw-r--r--tests/replication/test_pusher_shard.py2
-rw-r--r--tests/replication/test_sharded_event_persister.py8
8 files changed, 17 insertions, 17 deletions
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()