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)
|
