diff --git a/tests/storage/test_event_chain.py b/tests/storage/test_event_chain.py
index 8d7310d8e5..2a9aa9e21c 100644
--- a/tests/storage/test_event_chain.py
+++ b/tests/storage/test_event_chain.py
@@ -417,7 +417,6 @@ class EventChainStoreTestCase(HomeserverTestCase):
def fetch_chains(
self, events: List[EventBase]
) -> Tuple[Dict[str, Tuple[int, int]], _LinkMap]:
-
# Fetch the map from event ID -> (chain ID, sequence number)
rows = self.get_success(
self.store.db_pool.simple_select_many_batch(
@@ -492,7 +491,6 @@ class LinkMapTestCase(unittest.TestCase):
class EventChainBackgroundUpdateTestCase(HomeserverTestCase):
-
servlets = [
admin.register_servlets,
room.register_servlets,
@@ -524,7 +522,8 @@ class EventChainBackgroundUpdateTestCase(HomeserverTestCase):
latest_event_ids = self.get_success(
self.store.get_prev_events_for_room(room_id)
)
- event, context, _ = self.get_success(
+
+ event, unpersisted_context, _ = self.get_success(
event_handler.create_event(
self.requester,
{
@@ -537,6 +536,7 @@ class EventChainBackgroundUpdateTestCase(HomeserverTestCase):
prev_event_ids=latest_event_ids,
)
)
+ context = self.get_success(unpersisted_context.persist(event))
self.get_success(
event_handler.handle_new_client_event(
self.requester, events_and_context=[(event, context)]
@@ -546,7 +546,7 @@ class EventChainBackgroundUpdateTestCase(HomeserverTestCase):
assert state_ids1 is not None
state1 = set(state_ids1.values())
- event, context, _ = self.get_success(
+ event, unpersisted_context, _ = self.get_success(
event_handler.create_event(
self.requester,
{
@@ -559,6 +559,7 @@ class EventChainBackgroundUpdateTestCase(HomeserverTestCase):
prev_event_ids=latest_event_ids,
)
)
+ context = self.get_success(unpersisted_context.persist(event))
self.get_success(
event_handler.handle_new_client_event(
self.requester, events_and_context=[(event, context)]
|