diff --git a/tests/handlers/test_federation.py b/tests/handlers/test_federation.py
index ec00900621..500c9ccfbc 100644
--- a/tests/handlers/test_federation.py
+++ b/tests/handlers/test_federation.py
@@ -50,7 +50,7 @@ class FederationTestCase(unittest.FederatingHomeserverTestCase):
hs = self.setup_test_homeserver(federation_http_client=None)
self.handler = hs.get_federation_handler()
self.store = hs.get_datastores().main
- self.state_storage = hs.get_storage().state
+ self.state_storage_controller = hs.get_storage_controllers().state
self._event_auth_handler = hs.get_event_auth_handler()
return hs
@@ -338,7 +338,9 @@ class FederationTestCase(unittest.FederatingHomeserverTestCase):
# mapping from (type, state_key) -> state_event_id
assert most_recent_prev_event_id is not None
prev_state_map = self.get_success(
- self.state_storage.get_state_ids_for_event(most_recent_prev_event_id)
+ self.state_storage_controller.get_state_ids_for_event(
+ most_recent_prev_event_id
+ )
)
# List of state event ID's
prev_state_ids = list(prev_state_map.values())
diff --git a/tests/handlers/test_federation_event.py b/tests/handlers/test_federation_event.py
index e64b28f28b..1d5b2492c0 100644
--- a/tests/handlers/test_federation_event.py
+++ b/tests/handlers/test_federation_event.py
@@ -70,7 +70,7 @@ class FederationEventHandlerTests(unittest.FederatingHomeserverTestCase):
) -> None:
OTHER_USER = f"@user:{self.OTHER_SERVER_NAME}"
main_store = self.hs.get_datastores().main
- state_storage = self.hs.get_storage().state
+ state_storage_controller = self.hs.get_storage_controllers().state
# create the room
user_id = self.register_user("kermit", "test")
@@ -146,10 +146,11 @@ class FederationEventHandlerTests(unittest.FederatingHomeserverTestCase):
)
if prev_exists_as_outlier:
prev_event.internal_metadata.outlier = True
- persistence = self.hs.get_storage().persistence
+ persistence = self.hs.get_storage_controllers().persistence
self.get_success(
persistence.persist_event(
- prev_event, EventContext.for_outlier(self.hs.get_storage())
+ prev_event,
+ EventContext.for_outlier(self.hs.get_storage_controllers()),
)
)
else:
@@ -216,7 +217,7 @@ class FederationEventHandlerTests(unittest.FederatingHomeserverTestCase):
# check that the state at that event is as expected
state = self.get_success(
- state_storage.get_state_ids_for_event(pulled_event.event_id)
+ state_storage_controller.get_state_ids_for_event(pulled_event.event_id)
)
expected_state = {
(e.type, e.state_key): e.event_id for e in state_at_prev_event
diff --git a/tests/handlers/test_message.py b/tests/handlers/test_message.py
index f4f7ab4845..44da96c792 100644
--- a/tests/handlers/test_message.py
+++ b/tests/handlers/test_message.py
@@ -37,7 +37,9 @@ class EventCreationTestCase(unittest.HomeserverTestCase):
def prepare(self, reactor, clock, hs):
self.handler = self.hs.get_event_creation_handler()
- self.persist_event_storage = self.hs.get_storage().persistence
+ self._persist_event_storage_controller = (
+ self.hs.get_storage_controllers().persistence
+ )
self.user_id = self.register_user("tester", "foobar")
self.access_token = self.login("tester", "foobar")
@@ -65,7 +67,9 @@ class EventCreationTestCase(unittest.HomeserverTestCase):
)
)
self.get_success(
- self.persist_event_storage.persist_event(memberEvent, memberEventContext)
+ self._persist_event_storage_controller.persist_event(
+ memberEvent, memberEventContext
+ )
)
return memberEvent, memberEventContext
@@ -129,7 +133,7 @@ class EventCreationTestCase(unittest.HomeserverTestCase):
self.assertNotEqual(event1.event_id, event3.event_id)
ret_event3, event_pos3, _ = self.get_success(
- self.persist_event_storage.persist_event(event3, context)
+ self._persist_event_storage_controller.persist_event(event3, context)
)
# Assert that the returned values match those from the initial event
@@ -143,7 +147,7 @@ class EventCreationTestCase(unittest.HomeserverTestCase):
self.assertNotEqual(event1.event_id, event3.event_id)
events, _ = self.get_success(
- self.persist_event_storage.persist_events([(event3, context)])
+ self._persist_event_storage_controller.persist_events([(event3, context)])
)
ret_event4 = events[0]
@@ -166,7 +170,7 @@ class EventCreationTestCase(unittest.HomeserverTestCase):
self.assertNotEqual(event1.event_id, event2.event_id)
events, _ = self.get_success(
- self.persist_event_storage.persist_events(
+ self._persist_event_storage_controller.persist_events(
[(event1, context1), (event2, context2)]
)
)
diff --git a/tests/handlers/test_user_directory.py b/tests/handlers/test_user_directory.py
index 4d658d29ca..a68c2ffd45 100644
--- a/tests/handlers/test_user_directory.py
+++ b/tests/handlers/test_user_directory.py
@@ -954,7 +954,7 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
)
self.get_success(
- self.hs.get_storage().persistence.persist_event(event, context)
+ self.hs.get_storage_controllers().persistence.persist_event(event, context)
)
def test_local_user_leaving_room_remains_in_user_directory(self) -> None:
|