From 164f50f5f25a3204cce5fd2c8f196e9a9d4deb5d Mon Sep 17 00:00:00 2001
From: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>
Date: Mon, 18 May 2020 10:43:05 +0100
Subject: [PATCH] fix mypy for tests/replication (#7518)

---
 changelog.d/7518.misc                          |  1 +
 tests/replication/slave/storage/test_events.py | 16 +++++-----------
 tests/replication/tcp/test_commands.py         |  4 ++--
 tox.ini                                        |  2 +-
 4 files changed, 9 insertions(+), 14 deletions(-)
 create mode 100644 changelog.d/7518.misc

diff --git a/changelog.d/7518.misc b/changelog.d/7518.misc
new file mode 100644
index 0000000000..f6e143fe1c
--- /dev/null
+++ b/changelog.d/7518.misc
@@ -0,0 +1 @@
+Fix typing annotations in `tests.replication`.
diff --git a/tests/replication/slave/storage/test_events.py b/tests/replication/slave/storage/test_events.py
index 0fee8a71c4..1a88c7fb80 100644
--- a/tests/replication/slave/storage/test_events.py
+++ b/tests/replication/slave/storage/test_events.py
@@ -17,11 +17,12 @@ from canonicaljson import encode_canonical_json
 
 from synapse.api.room_versions import RoomVersions
 from synapse.events import FrozenEvent, _EventInternalMetadata, make_event_from_dict
-from synapse.events.snapshot import EventContext
 from synapse.handlers.room import RoomEventSource
 from synapse.replication.slave.storage.events import SlavedEventStore
 from synapse.storage.roommember import RoomsForUser
 
+from tests.server import FakeTransport
+
 from ._base import BaseSlavedStoreTestCase
 
 USER_ID = "@feeling:test"
@@ -240,6 +241,7 @@ class SlavedEventStoreTestCase(BaseSlavedStoreTestCase):
 
         # limit the replication rate
         repl_transport = self._server_transport
+        assert isinstance(repl_transport, FakeTransport)
         repl_transport.autoflush = False
 
         # build the join and message events and persist them in the same batch.
@@ -322,7 +324,6 @@ class SlavedEventStoreTestCase(BaseSlavedStoreTestCase):
         type="m.room.message",
         key=None,
         internal={},
-        state=None,
         depth=None,
         prev_events=[],
         auth_events=[],
@@ -362,15 +363,8 @@ class SlavedEventStoreTestCase(BaseSlavedStoreTestCase):
         event = make_event_from_dict(event_dict, internal_metadata_dict=internal)
 
         self.event_id += 1
-
-        if state is not None:
-            state_ids = {key: e.event_id for key, e in state.items()}
-            context = EventContext.with_state(
-                state_group=None, current_state_ids=state_ids, prev_state_ids=state_ids
-            )
-        else:
-            state_handler = self.hs.get_state_handler()
-            context = self.get_success(state_handler.compute_event_context(event))
+        state_handler = self.hs.get_state_handler()
+        context = self.get_success(state_handler.compute_event_context(event))
 
         self.master_store.add_push_actions_to_staging(
             event.event_id, {user_id: actions for user_id, actions in push_actions}
diff --git a/tests/replication/tcp/test_commands.py b/tests/replication/tcp/test_commands.py
index 7ddfd0a733..60c10a441a 100644
--- a/tests/replication/tcp/test_commands.py
+++ b/tests/replication/tcp/test_commands.py
@@ -30,7 +30,7 @@ class ParseCommandTestCase(TestCase):
     def test_parse_rdata(self):
         line = 'RDATA events master 6287863 ["ev", ["$eventid", "!roomid", "type", null, null, null]]'
         cmd = parse_command_from_line(line)
-        self.assertIsInstance(cmd, RdataCommand)
+        assert isinstance(cmd, RdataCommand)
         self.assertEqual(cmd.stream_name, "events")
         self.assertEqual(cmd.instance_name, "master")
         self.assertEqual(cmd.token, 6287863)
@@ -38,7 +38,7 @@ class ParseCommandTestCase(TestCase):
     def test_parse_rdata_batch(self):
         line = 'RDATA presence master batch ["@foo:example.com", "online"]'
         cmd = parse_command_from_line(line)
-        self.assertIsInstance(cmd, RdataCommand)
+        assert isinstance(cmd, RdataCommand)
         self.assertEqual(cmd.stream_name, "presence")
         self.assertEqual(cmd.instance_name, "master")
         self.assertIsNone(cmd.token)
diff --git a/tox.ini b/tox.ini
index 5a1fa610b6..3bb4d45e2a 100644
--- a/tox.ini
+++ b/tox.ini
@@ -207,7 +207,7 @@ commands = mypy \
             synapse/storage/util \
             synapse/streams \
             synapse/util/caches/stream_change_cache.py \
-            tests/replication/tcp/streams \
+            tests/replication \
             tests/test_utils \
             tests/rest/client/v2_alpha/test_auth.py \
             tests/util/test_stream_change_cache.py