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