summary refs log tree commit diff
path: root/tests/federation
diff options
context:
space:
mode:
Diffstat (limited to 'tests/federation')
-rw-r--r--tests/federation/test_complexity.py4
-rw-r--r--tests/federation/test_federation_catch_up.py26
-rw-r--r--tests/federation/test_federation_sender.py6
-rw-r--r--tests/federation/transport/test_knocking.py2
4 files changed, 20 insertions, 18 deletions
diff --git a/tests/federation/test_complexity.py b/tests/federation/test_complexity.py
index e40ef95874..9336181c96 100644
--- a/tests/federation/test_complexity.py
+++ b/tests/federation/test_complexity.py
@@ -55,7 +55,7 @@ class RoomComplexityTests(unittest.FederatingHomeserverTestCase):
         self.assertTrue(complexity > 0, complexity)
 
         # Artificially raise the complexity
-        store = self.hs.get_datastore()
+        store = self.hs.get_datastores().main
         store.get_current_state_event_counts = lambda x: make_awaitable(500 * 1.23)
 
         # Get the room complexity again -- make sure it's our artificial value
@@ -149,7 +149,7 @@ class RoomComplexityTests(unittest.FederatingHomeserverTestCase):
         )
 
         # Artificially raise the complexity
-        self.hs.get_datastore().get_current_state_event_counts = (
+        self.hs.get_datastores().main.get_current_state_event_counts = (
             lambda x: make_awaitable(600)
         )
 
diff --git a/tests/federation/test_federation_catch_up.py b/tests/federation/test_federation_catch_up.py
index f0aa8ed9db..2873b4d430 100644
--- a/tests/federation/test_federation_catch_up.py
+++ b/tests/federation/test_federation_catch_up.py
@@ -64,7 +64,7 @@ class FederationCatchUpTestCases(FederatingHomeserverTestCase):
             Dictionary of { event_id: str, stream_ordering: int }
         """
         event_id, stream_ordering = self.get_success(
-            self.hs.get_datastore().db_pool.execute(
+            self.hs.get_datastores().main.db_pool.execute(
                 "test:get_destination_rooms",
                 None,
                 """
@@ -125,7 +125,7 @@ class FederationCatchUpTestCases(FederatingHomeserverTestCase):
         self.pump()
 
         lsso_1 = self.get_success(
-            self.hs.get_datastore().get_destination_last_successful_stream_ordering(
+            self.hs.get_datastores().main.get_destination_last_successful_stream_ordering(
                 "host2"
             )
         )
@@ -141,7 +141,7 @@ class FederationCatchUpTestCases(FederatingHomeserverTestCase):
         event_id_2 = self.helper.send(room, "rabbits!", tok=u1_token)["event_id"]
 
         lsso_2 = self.get_success(
-            self.hs.get_datastore().get_destination_last_successful_stream_ordering(
+            self.hs.get_datastores().main.get_destination_last_successful_stream_ordering(
                 "host2"
             )
         )
@@ -216,7 +216,9 @@ class FederationCatchUpTestCases(FederatingHomeserverTestCase):
 
         # let's also clear any backoffs
         self.get_success(
-            self.hs.get_datastore().set_destination_retry_timings("host2", None, 0, 0)
+            self.hs.get_datastores().main.set_destination_retry_timings(
+                "host2", None, 0, 0
+            )
         )
 
         # bring the remote online and clear the received pdu list
@@ -296,13 +298,13 @@ class FederationCatchUpTestCases(FederatingHomeserverTestCase):
 
         # destination_rooms should already be populated, but let us pretend that we already
         # sent (successfully) up to and including event id 2
-        event_2 = self.get_success(self.hs.get_datastore().get_event(event_id_2))
+        event_2 = self.get_success(self.hs.get_datastores().main.get_event(event_id_2))
 
         # also fetch event 5 so we know its last_successful_stream_ordering later
-        event_5 = self.get_success(self.hs.get_datastore().get_event(event_id_5))
+        event_5 = self.get_success(self.hs.get_datastores().main.get_event(event_id_5))
 
         self.get_success(
-            self.hs.get_datastore().set_destination_last_successful_stream_ordering(
+            self.hs.get_datastores().main.set_destination_last_successful_stream_ordering(
                 "host2", event_2.internal_metadata.stream_ordering
             )
         )
@@ -359,7 +361,7 @@ class FederationCatchUpTestCases(FederatingHomeserverTestCase):
         # ASSERT:
         # - All servers are up to date so none should have outstanding catch-up
         outstanding_when_successful = self.get_success(
-            self.hs.get_datastore().get_catch_up_outstanding_destinations(None)
+            self.hs.get_datastores().main.get_catch_up_outstanding_destinations(None)
         )
         self.assertEqual(outstanding_when_successful, [])
 
@@ -370,7 +372,7 @@ class FederationCatchUpTestCases(FederatingHomeserverTestCase):
         # - Mark zzzerver as being backed-off from
         now = self.clock.time_msec()
         self.get_success(
-            self.hs.get_datastore().set_destination_retry_timings(
+            self.hs.get_datastores().main.set_destination_retry_timings(
                 "zzzerver", now, now, 24 * 60 * 60 * 1000  # retry in 1 day
             )
         )
@@ -382,14 +384,14 @@ class FederationCatchUpTestCases(FederatingHomeserverTestCase):
         # - all remotes are outstanding
         # - they are returned in batches of 25, in order
         outstanding_1 = self.get_success(
-            self.hs.get_datastore().get_catch_up_outstanding_destinations(None)
+            self.hs.get_datastores().main.get_catch_up_outstanding_destinations(None)
         )
 
         self.assertEqual(len(outstanding_1), 25)
         self.assertEqual(outstanding_1, server_names[0:25])
 
         outstanding_2 = self.get_success(
-            self.hs.get_datastore().get_catch_up_outstanding_destinations(
+            self.hs.get_datastores().main.get_catch_up_outstanding_destinations(
                 outstanding_1[-1]
             )
         )
@@ -457,7 +459,7 @@ class FederationCatchUpTestCases(FederatingHomeserverTestCase):
         )
 
         self.get_success(
-            self.hs.get_datastore().set_destination_last_successful_stream_ordering(
+            self.hs.get_datastores().main.set_destination_last_successful_stream_ordering(
                 "host2", event_1.internal_metadata.stream_ordering
             )
         )
diff --git a/tests/federation/test_federation_sender.py b/tests/federation/test_federation_sender.py
index b2376e2db9..60e0c31f43 100644
--- a/tests/federation/test_federation_sender.py
+++ b/tests/federation/test_federation_sender.py
@@ -176,7 +176,7 @@ class FederationSenderDevicesTestCases(HomeserverTestCase):
         def get_users_who_share_room_with_user(user_id):
             return defer.succeed({"@user2:host2"})
 
-        hs.get_datastore().get_users_who_share_room_with_user = (
+        hs.get_datastores().main.get_users_who_share_room_with_user = (
             get_users_who_share_room_with_user
         )
 
@@ -395,7 +395,7 @@ class FederationSenderDevicesTestCases(HomeserverTestCase):
         # run the prune job
         self.reactor.advance(10)
         self.get_success(
-            self.hs.get_datastore()._prune_old_outbound_device_pokes(prune_age=1)
+            self.hs.get_datastores().main._prune_old_outbound_device_pokes(prune_age=1)
         )
 
         # recover the server
@@ -445,7 +445,7 @@ class FederationSenderDevicesTestCases(HomeserverTestCase):
         # run the prune job
         self.reactor.advance(10)
         self.get_success(
-            self.hs.get_datastore()._prune_old_outbound_device_pokes(prune_age=1)
+            self.hs.get_datastores().main._prune_old_outbound_device_pokes(prune_age=1)
         )
 
         # recover the server
diff --git a/tests/federation/transport/test_knocking.py b/tests/federation/transport/test_knocking.py
index 686f42ab48..adf0535d97 100644
--- a/tests/federation/transport/test_knocking.py
+++ b/tests/federation/transport/test_knocking.py
@@ -198,7 +198,7 @@ class FederationKnockingTestCase(
     ]
 
     def prepare(self, reactor, clock, homeserver):
-        self.store = homeserver.get_datastore()
+        self.store = homeserver.get_datastores().main
 
         # We're not going to be properly signing events as our remote homeserver is fake,
         # therefore disable event signature checks.