1 files changed, 5 insertions, 4 deletions
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