Refcator EventContext to accept state during init

1 files changed, 27 insertions, 21 deletions

diff --git a/synapse/events/snapshot.py b/synapse/events/snapshot.py
index f83a1581a6..fbbe8dd490 100644
--- a/synapse/events/snapshot.py
+++ b/synapse/events/snapshot.py
@@ -60,22 +60,22 @@ class EventContext(object):
         "app_service",
     ]
 
-    def __init__(self):
+    def __init__(self, state_group, current_state_ids, prev_state_ids,
+                 prev_group=None, delta_ids=None):
         # The current state including the current event
-        self.current_state_ids = None
+        self.current_state_ids = current_state_ids
         # The current state excluding the current event
-        self.prev_state_ids = None
-        self.state_group = None
-
-        self.rejected = False
+        self.prev_state_ids = prev_state_ids
+        self.state_group = state_group
 
         # A previously persisted state group and a delta between that
         # and this state.
-        self.prev_group = None
-        self.delta_ids = None
+        self.prev_group = prev_group
+        self.delta_ids = delta_ids
 
-        self.prev_state_events = None
+        self.prev_state_events = []
 
+        self.rejected = False
         self.app_service = None
 
     def serialize(self, event):
@@ -123,27 +123,33 @@ class EventContext(object):
         Returns:
             EventContext
         """
-        context = EventContext()
-        context.state_group = input["state_group"]
-        context.rejected = input["rejected"]
-        context.prev_group = input["prev_group"]
-        context.delta_ids = _decode_state_dict(input["delta_ids"])
-        context.prev_state_events = input["prev_state_events"]
-
         # We use the state_group and prev_state_id stuff to pull the
         # current_state_ids out of the DB and construct prev_state_ids.
         prev_state_id = input["prev_state_id"]
         event_type = input["event_type"]
         event_state_key = input["event_state_key"]
 
-        context.current_state_ids = yield store.get_state_ids_for_group(
-            context.state_group,
+        state_group = input["state_group"]
+
+        current_state_ids = yield store.get_state_ids_for_group(
+            state_group,
         )
         if prev_state_id and event_state_key:
-            context.prev_state_ids = dict(context.current_state_ids)
-            context.prev_state_ids[(event_type, event_state_key)] = prev_state_id
+            prev_state_ids = dict(current_state_ids)
+            prev_state_ids[(event_type, event_state_key)] = prev_state_id
         else:
-            context.prev_state_ids = context.current_state_ids
+            prev_state_ids = current_state_ids
+
+        context = EventContext(
+            state_group=state_group,
+            current_state_ids=current_state_ids,
+            prev_state_ids=prev_state_ids,
+            prev_group=input["prev_group"],
+            delta_ids = _decode_state_dict(input["delta_ids"]),
+        )
+
+        context.rejected = input["rejected"]
+        context.prev_state_events = input["prev_state_events"]
 
         app_service_id = input["app_service_id"]
         if app_service_id: