diff options
Diffstat (limited to '')
| -rw-r--r-- | tests/test_state.py | 198 |
1 files changed, 137 insertions, 61 deletions
diff --git a/tests/test_state.py b/tests/test_state.py index 504530b49a..b20a26e1ff 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -11,7 +11,19 @@ # 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 Collection, Dict, List, Optional, cast +from typing import ( + Any, + Collection, + Dict, + Generator, + Iterable, + Iterator, + List, + Optional, + Set, + Tuple, + cast, +) from unittest.mock import Mock from twisted.internet import defer @@ -19,9 +31,11 @@ from twisted.internet import defer from synapse.api.auth import Auth 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 import EventBase, make_event_from_dict from synapse.events.snapshot import EventContext from synapse.state import StateHandler, StateResolutionHandler, _make_state_cache_entry +from synapse.types import MutableStateMap, StateMap +from synapse.types.state import StateFilter from synapse.util import Clock from synapse.util.macaroons import MacaroonGenerator @@ -33,14 +47,14 @@ _next_event_id = 1000 def create_event( - name=None, - type=None, - state_key=None, - depth=2, - event_id=None, - prev_events: Optional[List[str]] = None, - **kwargs, -): + name: Optional[str] = None, + type: Optional[str] = None, + state_key: Optional[str] = None, + depth: int = 2, + event_id: Optional[str] = None, + prev_events: Optional[List[Tuple[str, dict]]] = None, + **kwargs: Any, +) -> EventBase: global _next_event_id if not event_id: @@ -67,21 +81,21 @@ def create_event( d.update(kwargs) - event = make_event_from_dict(d) - - return event + return make_event_from_dict(d) class _DummyStore: - def __init__(self): - self._event_to_state_group = {} - self._group_to_state = {} + def __init__(self) -> None: + self._event_to_state_group: Dict[str, int] = {} + self._group_to_state: Dict[int, MutableStateMap[str]] = {} - self._event_id_to_event = {} + self._event_id_to_event: Dict[str, EventBase] = {} self._next_group = 1 - async def get_state_groups_ids(self, room_id, event_ids): + async def get_state_groups_ids( + self, room_id: str, event_ids: Collection[str] + ) -> Dict[int, MutableStateMap[str]]: groups = {} for event_id in event_ids: group = self._event_to_state_group.get(event_id) @@ -90,16 +104,25 @@ class _DummyStore: return groups - async def get_state_ids_for_group(self, state_group, state_filter=None): + async def get_state_ids_for_group( + self, state_group: int, state_filter: Optional[StateFilter] = None + ) -> MutableStateMap[str]: return self._group_to_state[state_group] async def store_state_group( - self, event_id, room_id, prev_group, delta_ids, current_state_ids - ): + self, + event_id: str, + room_id: str, + prev_group: Optional[int], + delta_ids: Optional[StateMap[str]], + current_state_ids: Optional[StateMap[str]], + ) -> int: state_group = self._next_group self._next_group += 1 if current_state_ids is None: + assert prev_group is not None + assert delta_ids is not None current_state_ids = dict(self._group_to_state[prev_group]) current_state_ids.update(delta_ids) @@ -107,7 +130,9 @@ class _DummyStore: return state_group - async def get_events(self, event_ids, **kwargs): + async def get_events( + self, event_ids: Collection[str], **kwargs: Any + ) -> Dict[str, EventBase]: return { e_id: self._event_id_to_event[e_id] for e_id in event_ids @@ -119,31 +144,36 @@ class _DummyStore: ) -> Dict[str, bool]: return {e: False for e in event_ids} - async def get_state_group_delta(self, name): + async def get_state_group_delta( + self, name: str + ) -> Tuple[Optional[int], Optional[StateMap[str]]]: return None, None - def register_events(self, events): + def register_events(self, events: Iterable[EventBase]) -> None: for e in events: self._event_id_to_event[e.event_id] = e - def register_event_context(self, event, context): + def register_event_context(self, event: EventBase, context: EventContext) -> None: + assert context.state_group is not None self._event_to_state_group[event.event_id] = context.state_group - def register_event_id_state_group(self, event_id, state_group): + def register_event_id_state_group(self, event_id: str, state_group: int) -> None: self._event_to_state_group[event_id] = state_group - async def get_room_version_id(self, room_id): + async def get_room_version_id(self, room_id: str) -> str: return RoomVersions.V1.identifier async def get_state_group_for_events( - self, event_ids, await_full_state: bool = True - ): + self, event_ids: Collection[str], await_full_state: bool = True + ) -> Dict[str, int]: res = {} for event in event_ids: res[event] = self._event_to_state_group[event] return res - async def get_state_for_groups(self, groups): + async def get_state_for_groups( + self, groups: Collection[int] + ) -> Dict[int, MutableStateMap[str]]: res = {} for group in groups: state = self._group_to_state[group] @@ -152,21 +182,21 @@ class _DummyStore: class DictObj(dict): - def __init__(self, **kwargs): + def __init__(self, **kwargs: Any) -> None: super().__init__(kwargs) self.__dict__ = self class Graph: - def __init__(self, nodes, edges): - events = {} - clobbered = set(events.keys()) + def __init__(self, nodes: Dict[str, DictObj], edges: Dict[str, List[str]]): + events: Dict[str, EventBase] = {} + clobbered: Set[str] = set() for event_id, fields in nodes.items(): refs = edges.get(event_id) if refs: clobbered.difference_update(refs) - prev_events = [(r, {}) for r in refs] + prev_events: List[Tuple[str, dict]] = [(r, {}) for r in refs] else: prev_events = [] @@ -177,15 +207,12 @@ class Graph: self._leaves = clobbered self._events = sorted(events.values(), key=lambda e: e.depth) - def walk(self): + def walk(self) -> Iterator[EventBase]: return iter(self._events) - def get_leaves(self): - return (self._events[i] for i in self._leaves) - class StateTestCase(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.dummy_store = _DummyStore() storage_controllers = Mock(main=self.dummy_store, state=self.dummy_store) hs = Mock( @@ -220,7 +247,7 @@ class StateTestCase(unittest.TestCase): self.event_id = 0 @defer.inlineCallbacks - def test_branch_no_conflict(self): + def test_branch_no_conflict(self) -> Generator[defer.Deferred, Any, None]: graph = Graph( nodes={ "START": DictObj( @@ -248,6 +275,7 @@ class StateTestCase(unittest.TestCase): ctx_c = context_store["C"] ctx_d = context_store["D"] + prev_state_ids: StateMap[str] prev_state_ids = yield defer.ensureDeferred(ctx_d.get_prev_state_ids()) self.assertEqual(2, len(prev_state_ids)) @@ -255,7 +283,9 @@ class StateTestCase(unittest.TestCase): self.assertEqual(ctx_d.state_group_before_event, ctx_d.state_group) @defer.inlineCallbacks - def test_branch_basic_conflict(self): + def test_branch_basic_conflict( + self, + ) -> Generator["defer.Deferred[object]", Any, None]: graph = Graph( nodes={ "START": DictObj( @@ -280,7 +310,7 @@ class StateTestCase(unittest.TestCase): self.dummy_store.register_events(graph.walk()) - context_store = {} + context_store: Dict[str, EventContext] = {} for event in graph.walk(): context = yield defer.ensureDeferred( @@ -294,6 +324,7 @@ class StateTestCase(unittest.TestCase): ctx_c = context_store["C"] ctx_d = context_store["D"] + prev_state_ids: StateMap[str] prev_state_ids = yield defer.ensureDeferred(ctx_d.get_prev_state_ids()) self.assertSetEqual({"START", "A", "C"}, set(prev_state_ids.values())) @@ -301,7 +332,9 @@ class StateTestCase(unittest.TestCase): self.assertEqual(ctx_d.state_group_before_event, ctx_d.state_group) @defer.inlineCallbacks - def test_branch_have_banned_conflict(self): + def test_branch_have_banned_conflict( + self, + ) -> Generator["defer.Deferred[object]", Any, None]: graph = Graph( nodes={ "START": DictObj( @@ -338,7 +371,7 @@ class StateTestCase(unittest.TestCase): self.dummy_store.register_events(graph.walk()) - context_store = {} + context_store: Dict[str, EventContext] = {} for event in graph.walk(): context = yield defer.ensureDeferred( @@ -353,13 +386,16 @@ class StateTestCase(unittest.TestCase): ctx_c = context_store["C"] ctx_e = context_store["E"] + prev_state_ids: StateMap[str] prev_state_ids = yield defer.ensureDeferred(ctx_e.get_prev_state_ids()) self.assertSetEqual({"START", "A", "B", "C"}, set(prev_state_ids.values())) self.assertEqual(ctx_c.state_group, ctx_e.state_group_before_event) self.assertEqual(ctx_e.state_group_before_event, ctx_e.state_group) @defer.inlineCallbacks - def test_branch_have_perms_conflict(self): + def test_branch_have_perms_conflict( + self, + ) -> Generator["defer.Deferred[object]", Any, None]: userid1 = "@user_id:example.com" userid2 = "@user_id2:example.com" @@ -413,7 +449,7 @@ class StateTestCase(unittest.TestCase): self.dummy_store.register_events(graph.walk()) - context_store = {} + context_store: Dict[str, EventContext] = {} for event in graph.walk(): context = yield defer.ensureDeferred( @@ -428,14 +464,17 @@ class StateTestCase(unittest.TestCase): ctx_b = context_store["B"] ctx_d = context_store["D"] + prev_state_ids: StateMap[str] prev_state_ids = yield defer.ensureDeferred(ctx_d.get_prev_state_ids()) self.assertSetEqual({"A1", "A2", "A3", "A5", "B"}, set(prev_state_ids.values())) self.assertEqual(ctx_b.state_group, ctx_d.state_group_before_event) self.assertEqual(ctx_d.state_group_before_event, ctx_d.state_group) - def _add_depths(self, nodes, edges): - def _get_depth(ev): + def _add_depths( + self, nodes: Dict[str, DictObj], edges: Dict[str, List[str]] + ) -> None: + def _get_depth(ev: str) -> int: node = nodes[ev] if "depth" not in node: prevs = edges[ev] @@ -447,7 +486,9 @@ class StateTestCase(unittest.TestCase): _get_depth(n) @defer.inlineCallbacks - def test_annotate_with_old_message(self): + def test_annotate_with_old_message( + self, + ) -> Generator["defer.Deferred[object]", Any, None]: event = create_event(type="test_message", name="event") old_state = [ @@ -456,6 +497,7 @@ class StateTestCase(unittest.TestCase): create_event(type="test2", state_key=""), ] + context: EventContext context = yield defer.ensureDeferred( self.state.compute_event_context( event, @@ -466,9 +508,11 @@ class StateTestCase(unittest.TestCase): ) ) + prev_state_ids: StateMap[str] prev_state_ids = yield defer.ensureDeferred(context.get_prev_state_ids()) self.assertCountEqual((e.event_id for e in old_state), prev_state_ids.values()) + current_state_ids: StateMap[str] current_state_ids = yield defer.ensureDeferred(context.get_current_state_ids()) self.assertCountEqual( (e.event_id for e in old_state), current_state_ids.values() @@ -478,7 +522,9 @@ class StateTestCase(unittest.TestCase): self.assertEqual(context.state_group_before_event, context.state_group) @defer.inlineCallbacks - def test_annotate_with_old_state(self): + def test_annotate_with_old_state( + self, + ) -> Generator["defer.Deferred[object]", Any, None]: event = create_event(type="state", state_key="", name="event") old_state = [ @@ -487,6 +533,7 @@ class StateTestCase(unittest.TestCase): create_event(type="test2", state_key=""), ] + context: EventContext context = yield defer.ensureDeferred( self.state.compute_event_context( event, @@ -497,9 +544,11 @@ class StateTestCase(unittest.TestCase): ) ) + prev_state_ids: StateMap[str] prev_state_ids = yield defer.ensureDeferred(context.get_prev_state_ids()) self.assertCountEqual((e.event_id for e in old_state), prev_state_ids.values()) + current_state_ids: StateMap[str] current_state_ids = yield defer.ensureDeferred(context.get_current_state_ids()) self.assertCountEqual( (e.event_id for e in old_state + [event]), current_state_ids.values() @@ -511,7 +560,9 @@ class StateTestCase(unittest.TestCase): self.assertEqual({("state", ""): event.event_id}, context.delta_ids) @defer.inlineCallbacks - def test_trivial_annotate_message(self): + def test_trivial_annotate_message( + self, + ) -> Generator["defer.Deferred[object]", Any, None]: prev_event_id = "prev_event_id" event = create_event( type="test_message", name="event2", prev_events=[(prev_event_id, {})] @@ -534,8 +585,10 @@ class StateTestCase(unittest.TestCase): ) self.dummy_store.register_event_id_state_group(prev_event_id, group_name) + context: EventContext context = yield defer.ensureDeferred(self.state.compute_event_context(event)) + current_state_ids: StateMap[str] current_state_ids = yield defer.ensureDeferred(context.get_current_state_ids()) self.assertEqual( @@ -545,7 +598,9 @@ class StateTestCase(unittest.TestCase): self.assertEqual(group_name, context.state_group) @defer.inlineCallbacks - def test_trivial_annotate_state(self): + def test_trivial_annotate_state( + self, + ) -> Generator["defer.Deferred[object]", Any, None]: prev_event_id = "prev_event_id" event = create_event( type="state", state_key="", name="event2", prev_events=[(prev_event_id, {})] @@ -568,8 +623,10 @@ class StateTestCase(unittest.TestCase): ) self.dummy_store.register_event_id_state_group(prev_event_id, group_name) + context: EventContext context = yield defer.ensureDeferred(self.state.compute_event_context(event)) + prev_state_ids: StateMap[str] prev_state_ids = yield defer.ensureDeferred(context.get_prev_state_ids()) self.assertEqual({e.event_id for e in old_state}, set(prev_state_ids.values())) @@ -577,7 +634,9 @@ class StateTestCase(unittest.TestCase): self.assertIsNotNone(context.state_group) @defer.inlineCallbacks - def test_resolve_message_conflict(self): + def test_resolve_message_conflict( + self, + ) -> Generator["defer.Deferred[Any]", Any, None]: prev_event_id1 = "event_id1" prev_event_id2 = "event_id2" event = create_event( @@ -605,10 +664,12 @@ class StateTestCase(unittest.TestCase): self.dummy_store.register_events(old_state_1) self.dummy_store.register_events(old_state_2) + context: EventContext context = yield self._get_context( event, prev_event_id1, old_state_1, prev_event_id2, old_state_2 ) + current_state_ids: StateMap[str] current_state_ids = yield defer.ensureDeferred(context.get_current_state_ids()) self.assertEqual(len(current_state_ids), 6) @@ -616,7 +677,9 @@ class StateTestCase(unittest.TestCase): self.assertIsNotNone(context.state_group) @defer.inlineCallbacks - def test_resolve_state_conflict(self): + def test_resolve_state_conflict( + self, + ) -> Generator["defer.Deferred[Any]", Any, None]: prev_event_id1 = "event_id1" prev_event_id2 = "event_id2" event = create_event( @@ -645,12 +708,14 @@ class StateTestCase(unittest.TestCase): store = _DummyStore() store.register_events(old_state_1) store.register_events(old_state_2) - self.dummy_store.get_events = store.get_events + self.dummy_store.get_events = store.get_events # type: ignore[assignment] + context: EventContext context = yield self._get_context( event, prev_event_id1, old_state_1, prev_event_id2, old_state_2 ) + current_state_ids: StateMap[str] current_state_ids = yield defer.ensureDeferred(context.get_current_state_ids()) self.assertEqual(len(current_state_ids), 6) @@ -658,7 +723,9 @@ class StateTestCase(unittest.TestCase): self.assertIsNotNone(context.state_group) @defer.inlineCallbacks - def test_standard_depth_conflict(self): + def test_standard_depth_conflict( + self, + ) -> Generator["defer.Deferred[Any]", Any, None]: prev_event_id1 = "event_id1" prev_event_id2 = "event_id2" event = create_event( @@ -700,12 +767,14 @@ class StateTestCase(unittest.TestCase): store = _DummyStore() store.register_events(old_state_1) store.register_events(old_state_2) - self.dummy_store.get_events = store.get_events + self.dummy_store.get_events = store.get_events # type: ignore[assignment] + context: EventContext context = yield self._get_context( event, prev_event_id1, old_state_1, prev_event_id2, old_state_2 ) + current_state_ids: StateMap[str] current_state_ids = yield defer.ensureDeferred(context.get_current_state_ids()) self.assertEqual(old_state_2[3].event_id, current_state_ids[("test1", "1")]) @@ -740,8 +809,14 @@ class StateTestCase(unittest.TestCase): @defer.inlineCallbacks def _get_context( - self, event, prev_event_id_1, old_state_1, prev_event_id_2, old_state_2 - ): + self, + event: EventBase, + prev_event_id_1: str, + old_state_1: Collection[EventBase], + prev_event_id_2: str, + old_state_2: Collection[EventBase], + ) -> Generator["defer.Deferred[object]", Any, EventContext]: + sg1: int sg1 = yield defer.ensureDeferred( self.dummy_store.store_state_group( prev_event_id_1, @@ -753,6 +828,7 @@ class StateTestCase(unittest.TestCase): ) self.dummy_store.register_event_id_state_group(prev_event_id_1, sg1) + sg2: int sg2 = yield defer.ensureDeferred( self.dummy_store.store_state_group( prev_event_id_2, @@ -767,7 +843,7 @@ class StateTestCase(unittest.TestCase): result = yield defer.ensureDeferred(self.state.compute_event_context(event)) return result - def test_make_state_cache_entry(self): + def test_make_state_cache_entry(self) -> None: "Test that calculating a prev_group and delta is correct" new_state = { |