1 files changed, 45 insertions, 1 deletions
diff --git a/tests/test_state.py b/tests/test_state.py
index 6ca8d8f21d..bafd6d1750 100644
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -21,7 +21,7 @@ from synapse.api.constants import EventTypes, Membership
 from synapse.api.room_versions import RoomVersions
 from synapse.events import make_event_from_dict
 from synapse.events.snapshot import EventContext
-from synapse.state import StateHandler, StateResolutionHandler
+from synapse.state import StateHandler, StateResolutionHandler, _make_state_cache_entry
 from synapse.util import Clock
 from synapse.util.macaroons import MacaroonGenerator
 
@@ -99,6 +99,10 @@ class _DummyStore:
         state_group = self._next_group
         self._next_group += 1
 
+        if current_state_ids is None:
+            current_state_ids = dict(self._group_to_state[prev_group])
+            current_state_ids.update(delta_ids)
+
         self._group_to_state[state_group] = dict(current_state_ids)
 
         return state_group
@@ -760,3 +764,43 @@ class StateTestCase(unittest.TestCase):
 
         result = yield defer.ensureDeferred(self.state.compute_event_context(event))
         return result
+
+    def test_make_state_cache_entry(self):
+        "Test that calculating a prev_group and delta is correct"
+
+        new_state = {
+            ("a", ""): "E",
+            ("b", ""): "E",
+            ("c", ""): "E",
+            ("d", ""): "E",
+        }
+
+        # old_state_1 has fewer differences to new_state than old_state_2, but
+        # the delta involves deleting a key, which isn't allowed in the deltas,
+        # so we should pick old_state_2 as the prev_group.
+
+        # `old_state_1` has two differences: `a` and `e`
+        old_state_1 = {
+            ("a", ""): "F",
+            ("b", ""): "E",
+            ("c", ""): "E",
+            ("d", ""): "E",
+            ("e", ""): "E",
+        }
+
+        # `old_state_2` has three differences: `a`, `c` and `d`
+        old_state_2 = {
+            ("a", ""): "F",
+            ("b", ""): "E",
+            ("c", ""): "F",
+            ("d", ""): "F",
+        }
+
+        entry = _make_state_cache_entry(new_state, {1: old_state_1, 2: old_state_2})
+
+        self.assertEqual(entry.prev_group, 2)
+
+        # There are three changes from `old_state_2` to `new_state`
+        self.assertEqual(
+            entry.delta_ids, {("a", ""): "E", ("c", ""): "E", ("d", ""): "E"}
+        )