diff --git a/tests/replication/slave/storage/test_events.py b/tests/replication/slave/storage/test_events.py
index 105e1228bb..f430cce931 100644
--- a/tests/replication/slave/storage/test_events.py
+++ b/tests/replication/slave/storage/test_events.py
@@ -226,11 +226,9 @@ class SlavedEventStoreTestCase(BaseSlavedStoreTestCase):
context = EventContext()
context.current_state_ids = state_ids
context.prev_state_ids = state_ids
- elif not backfill:
+ else:
state_handler = self.hs.get_state_handler()
context = yield state_handler.compute_event_context(event)
- else:
- context = EventContext()
context.push_actions = push_actions
diff --git a/tests/storage/test_redaction.py b/tests/storage/test_redaction.py
index 6afaca3a61..de6d7904e3 100644
--- a/tests/storage/test_redaction.py
+++ b/tests/storage/test_redaction.py
@@ -36,8 +36,7 @@ class RedactionTestCase(unittest.TestCase):
self.store = hs.get_datastore()
self.event_builder_factory = hs.get_event_builder_factory()
- self.handlers = hs.get_handlers()
- self.message_handler = self.handlers.message_handler
+ self.event_creation_handler = hs.get_event_creation_handler()
self.u_alice = UserID.from_string("@alice:test")
self.u_bob = UserID.from_string("@bob:test")
@@ -59,7 +58,7 @@ class RedactionTestCase(unittest.TestCase):
"content": content,
})
- event, context = yield self.message_handler._create_new_client_event(
+ event, context = yield self.event_creation_handler._create_new_client_event(
builder
)
@@ -79,7 +78,7 @@ class RedactionTestCase(unittest.TestCase):
"content": {"body": body, "msgtype": u"message"},
})
- event, context = yield self.message_handler._create_new_client_event(
+ event, context = yield self.event_creation_handler._create_new_client_event(
builder
)
@@ -98,7 +97,7 @@ class RedactionTestCase(unittest.TestCase):
"redacts": event_id,
})
- event, context = yield self.message_handler._create_new_client_event(
+ event, context = yield self.event_creation_handler._create_new_client_event(
builder
)
diff --git a/tests/storage/test_roommember.py b/tests/storage/test_roommember.py
index 1be7d932f6..4aff38bd58 100644
--- a/tests/storage/test_roommember.py
+++ b/tests/storage/test_roommember.py
@@ -37,8 +37,7 @@ class RoomMemberStoreTestCase(unittest.TestCase):
# storage logic
self.store = hs.get_datastore()
self.event_builder_factory = hs.get_event_builder_factory()
- self.handlers = hs.get_handlers()
- self.message_handler = self.handlers.message_handler
+ self.event_creation_handler = hs.get_event_creation_handler()
self.u_alice = UserID.from_string("@alice:test")
self.u_bob = UserID.from_string("@bob:test")
@@ -58,7 +57,7 @@ class RoomMemberStoreTestCase(unittest.TestCase):
"content": {"membership": membership},
})
- event, context = yield self.message_handler._create_new_client_event(
+ event, context = yield self.event_creation_handler._create_new_client_event(
builder
)
diff --git a/tests/test_state.py b/tests/test_state.py
index feb84f3d48..253aa62f2a 100644
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -99,6 +99,10 @@ class StateGroupStore(object):
for e in events:
self._event_id_to_event[e.event_id] = e
+ def store_state_group(self, *args, **kwargs):
+ self._next_group += 1
+ return self._next_group
+
class DictObj(dict):
def __init__(self, **kwargs):
@@ -144,6 +148,7 @@ class StateTestCase(unittest.TestCase):
"get_events",
"get_next_state_group",
"get_state_group_delta",
+ "store_state_group",
]
)
hs = Mock(spec_set=[
@@ -316,6 +321,7 @@ class StateTestCase(unittest.TestCase):
store = StateGroupStore()
self.store.get_state_groups_ids.side_effect = store.get_state_groups_ids
self.store.get_events = store.get_events
+ self.store.store_state_group = store.store_state_group
store.register_events(graph.walk())
context_store = {}
@@ -399,6 +405,7 @@ class StateTestCase(unittest.TestCase):
store = StateGroupStore()
self.store.get_state_groups_ids.side_effect = store.get_state_groups_ids
self.store.get_events = store.get_events
+ self.store.store_state_group = store.store_state_group
store.register_events(graph.walk())
context_store = {}
|