1 files changed, 19 insertions, 0 deletions
diff --git a/synapse/events/snapshot.py b/synapse/events/snapshot.py
index f9568638a1..b090751bf1 100644
--- a/synapse/events/snapshot.py
+++ b/synapse/events/snapshot.py
@@ -228,6 +228,25 @@ class EventContext(object):
         else:
             self._prev_state_ids = self._current_state_ids
 
+    @defer.inlineCallbacks
+    def update_state(self, state_group, prev_state_ids, current_state_ids,
+                     delta_ids):
+        """Replace the state in the context
+        """
+
+        # We need to make sure we wait for any ongoing fetching of state
+        # to complete so that the updated state doesn't get clobbered
+        if self._fetching_state_deferred:
+            yield make_deferred_yieldable(self._fetching_state_deferred)
+
+        self.state_group = state_group
+        self._prev_state_ids = prev_state_ids
+        self._current_state_ids = current_state_ids
+        self.delta_ids = delta_ids
+
+        # We need to ensure that that we've marked as having fetched the state
+        self._fetching_state_deferred = defer.succeed(None)
+
 
 def _encode_state_dict(state_dict):
     """Since dicts of (type, state_key) -> event_id cannot be serialized in