2 files changed, 7 insertions, 13 deletions
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)