Add missing types to test_state. (#14985)

3 files changed, 141 insertions, 62 deletions

diff --git a/changelog.d/14985.misc b/changelog.d/14985.misc
new file mode 100644
index 0000000000..93ceaeafc9
--- /dev/null
+++ b/changelog.d/14985.misc
@@ -0,0 +1 @@
+Improve type hints.
diff --git a/mypy.ini b/mypy.ini
index 2f8806e1f4..93de1c97ea 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -38,7 +38,6 @@ exclude = (?x)
    |tests/module_api/test_api.py
    |tests/rest/media/v1/test_media_storage.py
    |tests/server.py
-   |tests/test_state.py
    |tests/test_terms_auth.py
    )$
 
@@ -117,6 +116,9 @@ disallow_untyped_defs = True
 [mypy-tests.test_server]
 disallow_untyped_defs = True
 
+[mypy-tests.test_state]
+disallow_untyped_defs = True
+
 [mypy-tests.types.*]
 disallow_untyped_defs = True
 
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 = {