2 files changed, 130 insertions, 0 deletions
diff --git a/tests/storage/util/__init__.py b/tests/storage/util/__init__.py
new file mode 100644
index 0000000000..3a5f22c022
--- /dev/null
+++ b/tests/storage/util/__init__.py
@@ -0,0 +1,13 @@
+# Copyright 2022 The Matrix.org Foundation C.I.C.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
diff --git a/tests/storage/util/test_partial_state_events_tracker.py b/tests/storage/util/test_partial_state_events_tracker.py
new file mode 100644
index 0000000000..303e190b6c
--- /dev/null
+++ b/tests/storage/util/test_partial_state_events_tracker.py
@@ -0,0 +1,117 @@
+# Copyright 2022 The Matrix.org Foundation C.I.C.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict
+from unittest import mock
+
+from twisted.internet.defer import CancelledError, ensureDeferred
+
+from synapse.storage.util.partial_state_events_tracker import PartialStateEventsTracker
+
+from tests.unittest import TestCase
+
+
+class PartialStateEventsTrackerTestCase(TestCase):
+    def setUp(self) -> None:
+        # the results to be returned by the mocked get_partial_state_events
+        self._events_dict: Dict[str, bool] = {}
+
+        async def get_partial_state_events(events):
+            return {e: self._events_dict[e] for e in events}
+
+        self.mock_store = mock.Mock(spec_set=["get_partial_state_events"])
+        self.mock_store.get_partial_state_events.side_effect = get_partial_state_events
+
+        self.tracker = PartialStateEventsTracker(self.mock_store)
+
+    def test_does_not_block_for_full_state_events(self):
+        self._events_dict = {"event1": False, "event2": False}
+
+        self.successResultOf(
+            ensureDeferred(self.tracker.await_full_state(["event1", "event2"]))
+        )
+
+        self.mock_store.get_partial_state_events.assert_called_once_with(
+            ["event1", "event2"]
+        )
+
+    def test_blocks_for_partial_state_events(self):
+        self._events_dict = {"event1": True, "event2": False}
+
+        d = ensureDeferred(self.tracker.await_full_state(["event1", "event2"]))
+
+        # there should be no result yet
+        self.assertNoResult(d)
+
+        # notifying that the event has been de-partial-stated should unblock
+        self.tracker.notify_un_partial_stated("event1")
+        self.successResultOf(d)
+
+    def test_un_partial_state_race(self):
+        # if the event is un-partial-stated between the initial check and the
+        # registration of the listener, it should not block.
+        self._events_dict = {"event1": True, "event2": False}
+
+        async def get_partial_state_events(events):
+            res = {e: self._events_dict[e] for e in events}
+            # change the result for next time
+            self._events_dict = {"event1": False, "event2": False}
+            return res
+
+        self.mock_store.get_partial_state_events.side_effect = get_partial_state_events
+
+        self.successResultOf(
+            ensureDeferred(self.tracker.await_full_state(["event1", "event2"]))
+        )
+
+    def test_un_partial_state_during_get_partial_state_events(self):
+        # we should correctly handle a call to notify_un_partial_stated during the
+        # second call to get_partial_state_events.
+
+        self._events_dict = {"event1": True, "event2": False}
+
+        async def get_partial_state_events1(events):
+            self.mock_store.get_partial_state_events.side_effect = (
+                get_partial_state_events2
+            )
+            return {e: self._events_dict[e] for e in events}
+
+        async def get_partial_state_events2(events):
+            self.tracker.notify_un_partial_stated("event1")
+            self._events_dict["event1"] = False
+            return {e: self._events_dict[e] for e in events}
+
+        self.mock_store.get_partial_state_events.side_effect = get_partial_state_events1
+
+        self.successResultOf(
+            ensureDeferred(self.tracker.await_full_state(["event1", "event2"]))
+        )
+
+    def test_cancellation(self):
+        self._events_dict = {"event1": True, "event2": False}
+
+        d1 = ensureDeferred(self.tracker.await_full_state(["event1", "event2"]))
+        self.assertNoResult(d1)
+
+        d2 = ensureDeferred(self.tracker.await_full_state(["event1"]))
+        self.assertNoResult(d2)
+
+        d1.cancel()
+        self.assertFailure(d1, CancelledError)
+
+        # d2 should still be waiting!
+        self.assertNoResult(d2)
+
+        self.tracker.notify_un_partial_stated("event1")
+        self.successResultOf(d2)