New upstream version 1.22.1

49 files changed, 2162 insertions, 690 deletions

diff --git a/tests/api/test_auth.py b/tests/api/test_auth.py
index 8ab56ec9..cb6f29d6 100644
--- a/tests/api/test_auth.py
+++ b/tests/api/test_auth.py
@@ -19,7 +19,6 @@ import pymacaroons
 
 from twisted.internet import defer
 
-import synapse.handlers.auth
 from synapse.api.auth import Auth
 from synapse.api.constants import UserTypes
 from synapse.api.errors import (
@@ -36,20 +35,15 @@ from tests import unittest
 from tests.utils import mock_getRawHeaders, setup_test_homeserver
 
 
-class TestHandlers:
-    def __init__(self, hs):
-        self.auth_handler = synapse.handlers.auth.AuthHandler(hs)
-
-
 class AuthTestCase(unittest.TestCase):
     @defer.inlineCallbacks
     def setUp(self):
         self.state_handler = Mock()
         self.store = Mock()
 
-        self.hs = yield setup_test_homeserver(self.addCleanup, handlers=None)
+        self.hs = yield setup_test_homeserver(self.addCleanup)
         self.hs.get_datastore = Mock(return_value=self.store)
-        self.hs.handlers = TestHandlers(self.hs)
+        self.hs.get_auth_handler().store = self.store
         self.auth = Auth(self.hs)
 
         # AuthBlocking reads from the hs' config on initialization. We need to
@@ -283,7 +277,7 @@ class AuthTestCase(unittest.TestCase):
         self.store.get_device = Mock(return_value=defer.succeed(None))
 
         token = yield defer.ensureDeferred(
-            self.hs.handlers.auth_handler.get_access_token_for_user_id(
+            self.hs.get_auth_handler().get_access_token_for_user_id(
                 USER_ID, "DEVICE", valid_until_ms=None
             )
         )
diff --git a/tests/api/test_filtering.py b/tests/api/test_filtering.py
index d2d535d2..c98ae759 100644
--- a/tests/api/test_filtering.py
+++ b/tests/api/test_filtering.py
@@ -50,10 +50,7 @@ class FilteringTestCase(unittest.TestCase):
         self.mock_http_client.put_json = DeferredMockCallable()
 
         hs = yield setup_test_homeserver(
-            self.addCleanup,
-            handlers=None,
-            http_client=self.mock_http_client,
-            keyring=Mock(),
+            self.addCleanup, http_client=self.mock_http_client, keyring=Mock(),
         )
 
         self.filtering = hs.get_filtering()
diff --git a/tests/app/test_frontend_proxy.py b/tests/app/test_frontend_proxy.py
index 641093d3..4a301b84 100644
--- a/tests/app/test_frontend_proxy.py
+++ b/tests/app/test_frontend_proxy.py
@@ -22,7 +22,7 @@ class FrontendProxyTests(HomeserverTestCase):
     def make_homeserver(self, reactor, clock):
 
         hs = self.setup_test_homeserver(
-            http_client=None, homeserverToUse=GenericWorkerServer
+            http_client=None, homeserver_to_use=GenericWorkerServer
         )
 
         return hs
diff --git a/tests/app/test_openid_listener.py b/tests/app/test_openid_listener.py
index 0f016c32..c2b10d2c 100644
--- a/tests/app/test_openid_listener.py
+++ b/tests/app/test_openid_listener.py
@@ -26,7 +26,7 @@ from tests.unittest import HomeserverTestCase
 class FederationReaderOpenIDListenerTests(HomeserverTestCase):
     def make_homeserver(self, reactor, clock):
         hs = self.setup_test_homeserver(
-            http_client=None, homeserverToUse=GenericWorkerServer
+            http_client=None, homeserver_to_use=GenericWorkerServer
         )
         return hs
 
@@ -84,7 +84,7 @@ class FederationReaderOpenIDListenerTests(HomeserverTestCase):
 class SynapseHomeserverOpenIDListenerTests(HomeserverTestCase):
     def make_homeserver(self, reactor, clock):
         hs = self.setup_test_homeserver(
-            http_client=None, homeserverToUse=SynapseHomeServer
+            http_client=None, homeserver_to_use=SynapseHomeServer
         )
         return hs
 
diff --git a/tests/appservice/test_scheduler.py b/tests/appservice/test_scheduler.py
index 68a4caab..97f8cad0 100644
--- a/tests/appservice/test_scheduler.py
+++ b/tests/appservice/test_scheduler.py
@@ -60,7 +60,7 @@ class ApplicationServiceSchedulerTransactionCtrlTestCase(unittest.TestCase):
         self.successResultOf(defer.ensureDeferred(self.txnctrl.send(service, events)))
 
         self.store.create_appservice_txn.assert_called_once_with(
-            service=service, events=events  # txn made and saved
+            service=service, events=events, ephemeral=[]  # txn made and saved
         )
         self.assertEquals(0, len(self.txnctrl.recoverers))  # no recoverer made
         txn.complete.assert_called_once_with(self.store)  # txn completed
@@ -81,7 +81,7 @@ class ApplicationServiceSchedulerTransactionCtrlTestCase(unittest.TestCase):
         self.successResultOf(defer.ensureDeferred(self.txnctrl.send(service, events)))
 
         self.store.create_appservice_txn.assert_called_once_with(
-            service=service, events=events  # txn made and saved
+            service=service, events=events, ephemeral=[]  # txn made and saved
         )
         self.assertEquals(0, txn.send.call_count)  # txn not sent though
         self.assertEquals(0, txn.complete.call_count)  # or completed
@@ -106,7 +106,7 @@ class ApplicationServiceSchedulerTransactionCtrlTestCase(unittest.TestCase):
         self.successResultOf(defer.ensureDeferred(self.txnctrl.send(service, events)))
 
         self.store.create_appservice_txn.assert_called_once_with(
-            service=service, events=events
+            service=service, events=events, ephemeral=[]
         )
         self.assertEquals(1, self.recoverer_fn.call_count)  # recoverer made
         self.assertEquals(1, self.recoverer.recover.call_count)  # and invoked
@@ -202,26 +202,28 @@ class ApplicationServiceSchedulerQueuerTestCase(unittest.TestCase):
         # Expect the event to be sent immediately.
         service = Mock(id=4)
         event = Mock()
-        self.queuer.enqueue(service, event)
-        self.txn_ctrl.send.assert_called_once_with(service, [event])
+        self.queuer.enqueue_event(service, event)
+        self.txn_ctrl.send.assert_called_once_with(service, [event], [])
 
     def test_send_single_event_with_queue(self):
         d = defer.Deferred()
-        self.txn_ctrl.send = Mock(side_effect=lambda x, y: make_deferred_yieldable(d))
+        self.txn_ctrl.send = Mock(
+            side_effect=lambda x, y, z: make_deferred_yieldable(d)
+        )
         service = Mock(id=4)
         event = Mock(event_id="first")
         event2 = Mock(event_id="second")
         event3 = Mock(event_id="third")
         # Send an event and don't resolve it just yet.
-        self.queuer.enqueue(service, event)
+        self.queuer.enqueue_event(service, event)
         # Send more events: expect send() to NOT be called multiple times.
-        self.queuer.enqueue(service, event2)
-        self.queuer.enqueue(service, event3)
-        self.txn_ctrl.send.assert_called_with(service, [event])
+        self.queuer.enqueue_event(service, event2)
+        self.queuer.enqueue_event(service, event3)
+        self.txn_ctrl.send.assert_called_with(service, [event], [])
         self.assertEquals(1, self.txn_ctrl.send.call_count)
         # Resolve the send event: expect the queued events to be sent
         d.callback(service)
-        self.txn_ctrl.send.assert_called_with(service, [event2, event3])
+        self.txn_ctrl.send.assert_called_with(service, [event2, event3], [])
         self.assertEquals(2, self.txn_ctrl.send.call_count)
 
     def test_multiple_service_queues(self):
@@ -239,21 +241,99 @@ class ApplicationServiceSchedulerQueuerTestCase(unittest.TestCase):
 
         send_return_list = [srv_1_defer, srv_2_defer]
 
-        def do_send(x, y):
+        def do_send(x, y, z):
             return make_deferred_yieldable(send_return_list.pop(0))
 
         self.txn_ctrl.send = Mock(side_effect=do_send)
 
         # send events for different ASes and make sure they are sent
-        self.queuer.enqueue(srv1, srv_1_event)
-        self.queuer.enqueue(srv1, srv_1_event2)
-        self.txn_ctrl.send.assert_called_with(srv1, [srv_1_event])
-        self.queuer.enqueue(srv2, srv_2_event)
-        self.queuer.enqueue(srv2, srv_2_event2)
-        self.txn_ctrl.send.assert_called_with(srv2, [srv_2_event])
+        self.queuer.enqueue_event(srv1, srv_1_event)
+        self.queuer.enqueue_event(srv1, srv_1_event2)
+        self.txn_ctrl.send.assert_called_with(srv1, [srv_1_event], [])
+        self.queuer.enqueue_event(srv2, srv_2_event)
+        self.queuer.enqueue_event(srv2, srv_2_event2)
+        self.txn_ctrl.send.assert_called_with(srv2, [srv_2_event], [])
 
         # make sure callbacks for a service only send queued events for THAT
         # service
         srv_2_defer.callback(srv2)
-        self.txn_ctrl.send.assert_called_with(srv2, [srv_2_event2])
+        self.txn_ctrl.send.assert_called_with(srv2, [srv_2_event2], [])
+        self.assertEquals(3, self.txn_ctrl.send.call_count)
+
+    def test_send_large_txns(self):
+        srv_1_defer = defer.Deferred()
+        srv_2_defer = defer.Deferred()
+        send_return_list = [srv_1_defer, srv_2_defer]
+
+        def do_send(x, y, z):
+            return make_deferred_yieldable(send_return_list.pop(0))
+
+        self.txn_ctrl.send = Mock(side_effect=do_send)
+
+        service = Mock(id=4, name="service")
+        event_list = [Mock(name="event%i" % (i + 1)) for i in range(200)]
+        for event in event_list:
+            self.queuer.enqueue_event(service, event)
+
+        # Expect the first event to be sent immediately.
+        self.txn_ctrl.send.assert_called_with(service, [event_list[0]], [])
+        srv_1_defer.callback(service)
+        # Then send the next 100 events
+        self.txn_ctrl.send.assert_called_with(service, event_list[1:101], [])
+        srv_2_defer.callback(service)
+        # Then the final 99 events
+        self.txn_ctrl.send.assert_called_with(service, event_list[101:], [])
         self.assertEquals(3, self.txn_ctrl.send.call_count)
+
+    def test_send_single_ephemeral_no_queue(self):
+        # Expect the event to be sent immediately.
+        service = Mock(id=4, name="service")
+        event_list = [Mock(name="event")]
+        self.queuer.enqueue_ephemeral(service, event_list)
+        self.txn_ctrl.send.assert_called_once_with(service, [], event_list)
+
+    def test_send_multiple_ephemeral_no_queue(self):
+        # Expect the event to be sent immediately.
+        service = Mock(id=4, name="service")
+        event_list = [Mock(name="event1"), Mock(name="event2"), Mock(name="event3")]
+        self.queuer.enqueue_ephemeral(service, event_list)
+        self.txn_ctrl.send.assert_called_once_with(service, [], event_list)
+
+    def test_send_single_ephemeral_with_queue(self):
+        d = defer.Deferred()
+        self.txn_ctrl.send = Mock(
+            side_effect=lambda x, y, z: make_deferred_yieldable(d)
+        )
+        service = Mock(id=4)
+        event_list_1 = [Mock(event_id="event1"), Mock(event_id="event2")]
+        event_list_2 = [Mock(event_id="event3"), Mock(event_id="event4")]
+        event_list_3 = [Mock(event_id="event5"), Mock(event_id="event6")]
+
+        # Send an event and don't resolve it just yet.
+        self.queuer.enqueue_ephemeral(service, event_list_1)
+        # Send more events: expect send() to NOT be called multiple times.
+        self.queuer.enqueue_ephemeral(service, event_list_2)
+        self.queuer.enqueue_ephemeral(service, event_list_3)
+        self.txn_ctrl.send.assert_called_with(service, [], event_list_1)
+        self.assertEquals(1, self.txn_ctrl.send.call_count)
+        # Resolve txn_ctrl.send
+        d.callback(service)
+        # Expect the queued events to be sent
+        self.txn_ctrl.send.assert_called_with(service, [], event_list_2 + event_list_3)
+        self.assertEquals(2, self.txn_ctrl.send.call_count)
+
+    def test_send_large_txns_ephemeral(self):
+        d = defer.Deferred()
+        self.txn_ctrl.send = Mock(
+            side_effect=lambda x, y, z: make_deferred_yieldable(d)
+        )
+        # Expect the event to be sent immediately.
+        service = Mock(id=4, name="service")
+        first_chunk = [Mock(name="event%i" % (i + 1)) for i in range(100)]
+        second_chunk = [Mock(name="event%i" % (i + 101)) for i in range(50)]
+        event_list = first_chunk + second_chunk
+        self.queuer.enqueue_ephemeral(service, event_list)
+        self.txn_ctrl.send.assert_called_once_with(service, [], first_chunk)
+        d.callback(service)
+        self.txn_ctrl.send.assert_called_with(service, [], second_chunk)
+        self.assertEquals(2, self.txn_ctrl.send.call_count)
diff --git a/tests/crypto/test_keyring.py b/tests/crypto/test_keyring.py
index 8ff1460c..697916a0 100644
--- a/tests/crypto/test_keyring.py
+++ b/tests/crypto/test_keyring.py
@@ -315,7 +315,7 @@ class KeyringTestCase(unittest.HomeserverTestCase):
 class ServerKeyFetcherTestCase(unittest.HomeserverTestCase):
     def make_homeserver(self, reactor, clock):
         self.http_client = Mock()
-        hs = self.setup_test_homeserver(handlers=None, http_client=self.http_client)
+        hs = self.setup_test_homeserver(http_client=self.http_client)
         return hs
 
     def test_get_keys_from_server(self):
@@ -395,9 +395,7 @@ class PerspectivesKeyFetcherTestCase(unittest.HomeserverTestCase):
             }
         ]
 
-        return self.setup_test_homeserver(
-            handlers=None, http_client=self.http_client, config=config
-        )
+        return self.setup_test_homeserver(http_client=self.http_client, config=config)
 
     def build_perspectives_response(
         self, server_name: str, signing_key: SigningKey, valid_until_ts: int,
diff --git a/tests/handlers/test_admin.py b/tests/handlers/test_admin.py
index fc37c432..5c2b4de1 100644
--- a/tests/handlers/test_admin.py
+++ b/tests/handlers/test_admin.py
@@ -35,7 +35,7 @@ class ExfiltrateData(unittest.HomeserverTestCase):
     ]
 
     def prepare(self, reactor, clock, hs):
-        self.admin_handler = hs.get_handlers().admin_handler
+        self.admin_handler = hs.get_admin_handler()
 
         self.user1 = self.register_user("user1", "password")
         self.token1 = self.login("user1", "password")
diff --git a/tests/handlers/test_appservice.py b/tests/handlers/test_appservice.py
index 2a0b7c1b..ee4f3da3 100644
--- a/tests/handlers/test_appservice.py
+++ b/tests/handlers/test_appservice.py
@@ -18,6 +18,7 @@ from mock import Mock
 from twisted.internet import defer
 
 from synapse.handlers.appservice import ApplicationServicesHandler
+from synapse.types import RoomStreamToken
 
 from tests.test_utils import make_awaitable
 from tests.utils import MockClock
@@ -61,7 +62,9 @@ class AppServiceHandlerTestCase(unittest.TestCase):
             defer.succeed((0, [event])),
             defer.succeed((0, [])),
         ]
-        yield defer.ensureDeferred(self.handler.notify_interested_services(0))
+        yield defer.ensureDeferred(
+            self.handler.notify_interested_services(RoomStreamToken(None, 0))
+        )
         self.mock_scheduler.submit_event_for_as.assert_called_once_with(
             interested_service, event
         )
@@ -80,7 +83,9 @@ class AppServiceHandlerTestCase(unittest.TestCase):
             defer.succeed((0, [event])),
             defer.succeed((0, [])),
         ]
-        yield defer.ensureDeferred(self.handler.notify_interested_services(0))
+        yield defer.ensureDeferred(
+            self.handler.notify_interested_services(RoomStreamToken(None, 0))
+        )
         self.mock_as_api.query_user.assert_called_once_with(services[0], user_id)
 
     @defer.inlineCallbacks
@@ -97,7 +102,9 @@ class AppServiceHandlerTestCase(unittest.TestCase):
             defer.succeed((0, [event])),
             defer.succeed((0, [])),
         ]
-        yield defer.ensureDeferred(self.handler.notify_interested_services(0))
+        yield defer.ensureDeferred(
+            self.handler.notify_interested_services(RoomStreamToken(None, 0))
+        )
         self.assertFalse(
             self.mock_as_api.query_user.called,
             "query_user called when it shouldn't have been.",
diff --git a/tests/handlers/test_auth.py b/tests/handlers/test_auth.py
index 97877c2e..b5055e01 100644
--- a/tests/handlers/test_auth.py
+++ b/tests/handlers/test_auth.py
@@ -21,24 +21,17 @@ from twisted.internet import defer
 import synapse
 import synapse.api.errors
 from synapse.api.errors import ResourceLimitError
-from synapse.handlers.auth import AuthHandler
 
 from tests import unittest
 from tests.test_utils import make_awaitable
 from tests.utils import setup_test_homeserver
 
 
-class AuthHandlers:
-    def __init__(self, hs):
-        self.auth_handler = AuthHandler(hs)
-
-
 class AuthTestCase(unittest.TestCase):
     @defer.inlineCallbacks
     def setUp(self):
-        self.hs = yield setup_test_homeserver(self.addCleanup, handlers=None)
-        self.hs.handlers = AuthHandlers(self.hs)
-        self.auth_handler = self.hs.handlers.auth_handler
+        self.hs = yield setup_test_homeserver(self.addCleanup)
+        self.auth_handler = self.hs.get_auth_handler()
         self.macaroon_generator = self.hs.get_macaroon_generator()
 
         # MAU tests
diff --git a/tests/handlers/test_device.py b/tests/handlers/test_device.py
index 969d44c7..4512c513 100644
--- a/tests/handlers/test_device.py
+++ b/tests/handlers/test_device.py
@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 # Copyright 2016 OpenMarket Ltd
 # Copyright 2018 New Vector Ltd
+# Copyright 2020 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.
@@ -224,3 +225,84 @@ class DeviceTestCase(unittest.HomeserverTestCase):
                 )
             )
             self.reactor.advance(1000)
+
+
+class DehydrationTestCase(unittest.HomeserverTestCase):
+    def make_homeserver(self, reactor, clock):
+        hs = self.setup_test_homeserver("server", http_client=None)
+        self.handler = hs.get_device_handler()
+        self.registration = hs.get_registration_handler()
+        self.auth = hs.get_auth()
+        self.store = hs.get_datastore()
+        return hs
+
+    def test_dehydrate_and_rehydrate_device(self):
+        user_id = "@boris:dehydration"
+
+        self.get_success(self.store.register_user(user_id, "foobar"))
+
+        # First check if we can store and fetch a dehydrated device
+        stored_dehydrated_device_id = self.get_success(
+            self.handler.store_dehydrated_device(
+                user_id=user_id,
+                device_data={"device_data": {"foo": "bar"}},
+                initial_device_display_name="dehydrated device",
+            )
+        )
+
+        retrieved_device_id, device_data = self.get_success(
+            self.handler.get_dehydrated_device(user_id=user_id)
+        )
+
+        self.assertEqual(retrieved_device_id, stored_dehydrated_device_id)
+        self.assertEqual(device_data, {"device_data": {"foo": "bar"}})
+
+        # Create a new login for the user and dehydrated the device
+        device_id, access_token = self.get_success(
+            self.registration.register_device(
+                user_id=user_id, device_id=None, initial_display_name="new device",
+            )
+        )
+
+        # Trying to claim a nonexistent device should throw an error
+        self.get_failure(
+            self.handler.rehydrate_device(
+                user_id=user_id,
+                access_token=access_token,
+                device_id="not the right device ID",
+            ),
+            synapse.api.errors.NotFoundError,
+        )
+
+        # dehydrating the right devices should succeed and change our device ID
+        # to the dehydrated device's ID
+        res = self.get_success(
+            self.handler.rehydrate_device(
+                user_id=user_id,
+                access_token=access_token,
+                device_id=retrieved_device_id,
+            )
+        )
+
+        self.assertEqual(res, {"success": True})
+
+        # make sure that our device ID has changed
+        user_info = self.get_success(self.auth.get_user_by_access_token(access_token))
+
+        self.assertEqual(user_info["device_id"], retrieved_device_id)
+
+        # make sure the device has the display name that was set from the login
+        res = self.get_success(self.handler.get_device(user_id, retrieved_device_id))
+
+        self.assertEqual(res["display_name"], "new device")
+
+        # make sure that the device ID that we were initially assigned no longer exists
+        self.get_failure(
+            self.handler.get_device(user_id, device_id),
+            synapse.api.errors.NotFoundError,
+        )
+
+        # make sure that there's no device available for dehydrating now
+        ret = self.get_success(self.handler.get_dehydrated_device(user_id=user_id))
+
+        self.assertIsNone(ret)
diff --git a/tests/handlers/test_directory.py b/tests/handlers/test_directory.py
index bc0c5aef..2ce6dc95 100644
--- a/tests/handlers/test_directory.py
+++ b/tests/handlers/test_directory.py
@@ -48,7 +48,7 @@ class DirectoryTestCase(unittest.HomeserverTestCase):
             federation_registry=self.mock_registry,
         )
 
-        self.handler = hs.get_handlers().directory_handler
+        self.handler = hs.get_directory_handler()
 
         self.store = hs.get_datastore()
 
@@ -110,7 +110,7 @@ class TestCreateAlias(unittest.HomeserverTestCase):
     ]
 
     def prepare(self, reactor, clock, hs):
-        self.handler = hs.get_handlers().directory_handler
+        self.handler = hs.get_directory_handler()
 
         # Create user
         self.admin_user = self.register_user("admin", "pass", admin=True)
@@ -173,7 +173,7 @@ class TestDeleteAlias(unittest.HomeserverTestCase):
 
     def prepare(self, reactor, clock, hs):
         self.store = hs.get_datastore()
-        self.handler = hs.get_handlers().directory_handler
+        self.handler = hs.get_directory_handler()
         self.state_handler = hs.get_state_handler()
 
         # Create user
@@ -289,7 +289,7 @@ class CanonicalAliasTestCase(unittest.HomeserverTestCase):
 
     def prepare(self, reactor, clock, hs):
         self.store = hs.get_datastore()
-        self.handler = hs.get_handlers().directory_handler
+        self.handler = hs.get_directory_handler()
         self.state_handler = hs.get_state_handler()
 
         # Create user
@@ -442,7 +442,7 @@ class TestRoomListSearchDisabled(unittest.HomeserverTestCase):
         self.assertEquals(200, channel.code, channel.result)
 
         self.room_list_handler = hs.get_room_list_handler()
-        self.directory_handler = hs.get_handlers().directory_handler
+        self.directory_handler = hs.get_directory_handler()
 
         return hs
 
diff --git a/tests/handlers/test_e2e_keys.py b/tests/handlers/test_e2e_keys.py
index 366dcfb6..924f29f0 100644
--- a/tests/handlers/test_e2e_keys.py
+++ b/tests/handlers/test_e2e_keys.py
@@ -33,13 +33,15 @@ class E2eKeysHandlerTestCase(unittest.TestCase):
         super().__init__(*args, **kwargs)
         self.hs = None  # type: synapse.server.HomeServer
         self.handler = None  # type: synapse.handlers.e2e_keys.E2eKeysHandler
+        self.store = None  # type: synapse.storage.Storage
 
     @defer.inlineCallbacks
     def setUp(self):
         self.hs = yield utils.setup_test_homeserver(
-            self.addCleanup, handlers=None, federation_client=mock.Mock()
+            self.addCleanup, federation_client=mock.Mock()
         )
         self.handler = synapse.handlers.e2e_keys.E2eKeysHandler(self.hs)
+        self.store = self.hs.get_datastore()
 
     @defer.inlineCallbacks
     def test_query_local_devices_no_devices(self):
@@ -172,6 +174,89 @@ class E2eKeysHandlerTestCase(unittest.TestCase):
         )
 
     @defer.inlineCallbacks
+    def test_fallback_key(self):
+        local_user = "@boris:" + self.hs.hostname
+        device_id = "xyz"
+        fallback_key = {"alg1:k1": "key1"}
+        otk = {"alg1:k2": "key2"}
+
+        # we shouldn't have any unused fallback keys yet
+        res = yield defer.ensureDeferred(
+            self.store.get_e2e_unused_fallback_key_types(local_user, device_id)
+        )
+        self.assertEqual(res, [])
+
+        yield defer.ensureDeferred(
+            self.handler.upload_keys_for_user(
+                local_user,
+                device_id,
+                {"org.matrix.msc2732.fallback_keys": fallback_key},
+            )
+        )
+
+        # we should now have an unused alg1 key
+        res = yield defer.ensureDeferred(
+            self.store.get_e2e_unused_fallback_key_types(local_user, device_id)
+        )
+        self.assertEqual(res, ["alg1"])
+
+        # claiming an OTK when no OTKs are available should return the fallback
+        # key
+        res = yield defer.ensureDeferred(
+            self.handler.claim_one_time_keys(
+                {"one_time_keys": {local_user: {device_id: "alg1"}}}, timeout=None
+            )
+        )
+        self.assertEqual(
+            res,
+            {"failures": {}, "one_time_keys": {local_user: {device_id: fallback_key}}},
+        )
+
+        # we shouldn't have any unused fallback keys again
+        res = yield defer.ensureDeferred(
+            self.store.get_e2e_unused_fallback_key_types(local_user, device_id)
+        )
+        self.assertEqual(res, [])
+
+        # claiming an OTK again should return the same fallback key
+        res = yield defer.ensureDeferred(
+            self.handler.claim_one_time_keys(
+                {"one_time_keys": {local_user: {device_id: "alg1"}}}, timeout=None
+            )
+        )
+        self.assertEqual(
+            res,
+            {"failures": {}, "one_time_keys": {local_user: {device_id: fallback_key}}},
+        )
+
+        # if the user uploads a one-time key, the next claim should fetch the
+        # one-time key, and then go back to the fallback
+        yield defer.ensureDeferred(
+            self.handler.upload_keys_for_user(
+                local_user, device_id, {"one_time_keys": otk}
+            )
+        )
+
+        res = yield defer.ensureDeferred(
+            self.handler.claim_one_time_keys(
+                {"one_time_keys": {local_user: {device_id: "alg1"}}}, timeout=None
+            )
+        )
+        self.assertEqual(
+            res, {"failures": {}, "one_time_keys": {local_user: {device_id: otk}}},
+        )
+
+        res = yield defer.ensureDeferred(
+            self.handler.claim_one_time_keys(
+                {"one_time_keys": {local_user: {device_id: "alg1"}}}, timeout=None
+            )
+        )
+        self.assertEqual(
+            res,
+            {"failures": {}, "one_time_keys": {local_user: {device_id: fallback_key}}},
+        )
+
+    @defer.inlineCallbacks
     def test_replace_master_key(self):
         """uploading a new signing key should make the old signing key unavailable"""
         local_user = "@boris:" + self.hs.hostname
diff --git a/tests/handlers/test_e2e_room_keys.py b/tests/handlers/test_e2e_room_keys.py
index 7adde9b9..45f201a3 100644
--- a/tests/handlers/test_e2e_room_keys.py
+++ b/tests/handlers/test_e2e_room_keys.py
@@ -54,7 +54,7 @@ class E2eRoomKeysHandlerTestCase(unittest.TestCase):
     @defer.inlineCallbacks
     def setUp(self):
         self.hs = yield utils.setup_test_homeserver(
-            self.addCleanup, handlers=None, replication_layer=mock.Mock()
+            self.addCleanup, replication_layer=mock.Mock()
         )
         self.handler = synapse.handlers.e2e_room_keys.E2eRoomKeysHandler(self.hs)
         self.local_user = "@boris:" + self.hs.hostname
diff --git a/tests/handlers/test_federation.py b/tests/handlers/test_federation.py
index 96fea586..9ef80fe5 100644
--- a/tests/handlers/test_federation.py
+++ b/tests/handlers/test_federation.py
@@ -38,7 +38,7 @@ class FederationTestCase(unittest.HomeserverTestCase):
 
     def make_homeserver(self, reactor, clock):
         hs = self.setup_test_homeserver(http_client=None)
-        self.handler = hs.get_handlers().federation_handler
+        self.handler = hs.get_federation_handler()
         self.store = hs.get_datastore()
         return hs
 
diff --git a/tests/handlers/test_message.py b/tests/handlers/test_message.py
new file mode 100644
index 00000000..9f6f21a6
--- /dev/null
+++ b/tests/handlers/test_message.py
@@ -0,0 +1,156 @@
+# -*- coding: utf-8 -*-
+# Copyright 2020 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.
+import logging
+from typing import Tuple
+
+from synapse.api.constants import EventTypes
+from synapse.events import EventBase
+from synapse.events.snapshot import EventContext
+from synapse.rest import admin
+from synapse.rest.client.v1 import login, room
+from synapse.types import create_requester
+from synapse.util.stringutils import random_string
+
+from tests import unittest
+
+logger = logging.getLogger(__name__)
+
+
+class EventCreationTestCase(unittest.HomeserverTestCase):
+    servlets = [
+        admin.register_servlets,
+        login.register_servlets,
+        room.register_servlets,
+    ]
+
+    def prepare(self, reactor, clock, hs):
+        self.handler = self.hs.get_event_creation_handler()
+        self.persist_event_storage = self.hs.get_storage().persistence
+
+        self.user_id = self.register_user("tester", "foobar")
+        self.access_token = self.login("tester", "foobar")
+        self.room_id = self.helper.create_room_as(self.user_id, tok=self.access_token)
+
+        self.info = self.get_success(
+            self.hs.get_datastore().get_user_by_access_token(self.access_token,)
+        )
+        self.token_id = self.info["token_id"]
+
+        self.requester = create_requester(self.user_id, access_token_id=self.token_id)
+
+    def _create_duplicate_event(self, txn_id: str) -> Tuple[EventBase, EventContext]:
+        """Create a new event with the given transaction ID. All events produced
+        by this method will be considered duplicates.
+        """
+
+        # We create a new event with a random body, as otherwise we'll produce
+        # *exactly* the same event with the same hash, and so same event ID.
+        return self.get_success(
+            self.handler.create_event(
+                self.requester,
+                {
+                    "type": EventTypes.Message,
+                    "room_id": self.room_id,
+                    "sender": self.requester.user.to_string(),
+                    "content": {"msgtype": "m.text", "body": random_string(5)},
+                },
+                txn_id=txn_id,
+            )
+        )
+
+    def test_duplicated_txn_id(self):
+        """Test that attempting to handle/persist an event with a transaction ID
+        that has already been persisted correctly returns the old event and does
+        *not* produce duplicate messages.
+        """
+
+        txn_id = "something_suitably_random"
+
+        event1, context = self._create_duplicate_event(txn_id)
+
+        ret_event1 = self.get_success(
+            self.handler.handle_new_client_event(self.requester, event1, context)
+        )
+        stream_id1 = ret_event1.internal_metadata.stream_ordering
+
+        self.assertEqual(event1.event_id, ret_event1.event_id)
+
+        event2, context = self._create_duplicate_event(txn_id)
+
+        # We want to test that the deduplication at the persit event end works,
+        # so we want to make sure we test with different events.
+        self.assertNotEqual(event1.event_id, event2.event_id)
+
+        ret_event2 = self.get_success(
+            self.handler.handle_new_client_event(self.requester, event2, context)
+        )
+        stream_id2 = ret_event2.internal_metadata.stream_ordering
+
+        # Assert that the returned values match those from the initial event
+        # rather than the new one.
+        self.assertEqual(ret_event1.event_id, ret_event2.event_id)
+        self.assertEqual(stream_id1, stream_id2)
+
+        # Let's test that calling `persist_event` directly also does the right
+        # thing.
+        event3, context = self._create_duplicate_event(txn_id)
+        self.assertNotEqual(event1.event_id, event3.event_id)
+
+        ret_event3, event_pos3, _ = self.get_success(
+            self.persist_event_storage.persist_event(event3, context)
+        )
+
+        # Assert that the returned values match those from the initial event
+        # rather than the new one.
+        self.assertEqual(ret_event1.event_id, ret_event3.event_id)
+        self.assertEqual(stream_id1, event_pos3.stream)
+
+        # Let's test that calling `persist_events` directly also does the right
+        # thing.
+        event4, context = self._create_duplicate_event(txn_id)
+        self.assertNotEqual(event1.event_id, event3.event_id)
+
+        events, _ = self.get_success(
+            self.persist_event_storage.persist_events([(event3, context)])
+        )
+        ret_event4 = events[0]
+
+        # Assert that the returned values match those from the initial event
+        # rather than the new one.
+        self.assertEqual(ret_event1.event_id, ret_event4.event_id)
+
+    def test_duplicated_txn_id_one_call(self):
+        """Test that we correctly handle duplicates that we try and persist at
+        the same time.
+        """
+
+        txn_id = "something_else_suitably_random"
+
+        # Create two duplicate events to persist at the same time
+        event1, context1 = self._create_duplicate_event(txn_id)
+        event2, context2 = self._create_duplicate_event(txn_id)
+
+        # Ensure their event IDs are different to start with
+        self.assertNotEqual(event1.event_id, event2.event_id)
+
+        events, _ = self.get_success(
+            self.persist_event_storage.persist_events(
+                [(event1, context1), (event2, context2)]
+            )
+        )
+
+        # Check that we've deduplicated the events.
+        self.assertEqual(len(events), 2)
+        self.assertEqual(events[0].event_id, events[1].event_id)
diff --git a/tests/handlers/test_oidc.py b/tests/handlers/test_oidc.py
index d5087e58..b6f436c0 100644
--- a/tests/handlers/test_oidc.py
+++ b/tests/handlers/test_oidc.py
@@ -286,9 +286,15 @@ class OidcHandlerTestCase(HomeserverTestCase):
                 h._validate_metadata,
             )
 
-        # Tests for configs that the userinfo endpoint
+        # Tests for configs that require the userinfo endpoint
         self.assertFalse(h._uses_userinfo)
-        h._scopes = []  # do not request the openid scope
+        self.assertEqual(h._user_profile_method, "auto")
+        h._user_profile_method = "userinfo_endpoint"
+        self.assertTrue(h._uses_userinfo)
+
+        # Revert the profile method and do not request the "openid" scope.
+        h._user_profile_method = "auto"
+        h._scopes = []
         self.assertTrue(h._uses_userinfo)
         self.assertRaisesRegex(ValueError, "userinfo_endpoint", h._validate_metadata)
 
diff --git a/tests/handlers/test_presence.py b/tests/handlers/test_presence.py
index 306dcfe9..8ed67640 100644
--- a/tests/handlers/test_presence.py
+++ b/tests/handlers/test_presence.py
@@ -470,7 +470,7 @@ class PresenceJoinTestCase(unittest.HomeserverTestCase):
     def prepare(self, reactor, clock, hs):
         self.federation_sender = hs.get_federation_sender()
         self.event_builder_factory = hs.get_event_builder_factory()
-        self.federation_handler = hs.get_handlers().federation_handler
+        self.federation_handler = hs.get_federation_handler()
         self.presence_handler = hs.get_presence_handler()
 
         # self.event_builder_for_2 = EventBuilderFactory(hs)
@@ -615,7 +615,7 @@ class PresenceJoinTestCase(unittest.HomeserverTestCase):
             self.store.get_latest_event_ids_in_room(room_id)
         )
 
-        event = self.get_success(builder.build(prev_event_ids))
+        event = self.get_success(builder.build(prev_event_ids, None))
 
         self.get_success(self.federation_handler.on_receive_pdu(hostname, event))
 
diff --git a/tests/handlers/test_profile.py b/tests/handlers/test_profile.py
index 8e95e53d..a69fa28b 100644
--- a/tests/handlers/test_profile.py
+++ b/tests/handlers/test_profile.py
@@ -20,7 +20,6 @@ from twisted.internet import defer
 
 import synapse.types
 from synapse.api.errors import AuthError, SynapseError
-from synapse.handlers.profile import MasterProfileHandler
 from synapse.types import UserID
 
 from tests import unittest
@@ -28,11 +27,6 @@ from tests.test_utils import make_awaitable
 from tests.utils import setup_test_homeserver
 
 
-class ProfileHandlers:
-    def __init__(self, hs):
-        self.profile_handler = MasterProfileHandler(hs)
-
-
 class ProfileTestCase(unittest.TestCase):
     """ Tests profile management. """
 
@@ -51,7 +45,6 @@ class ProfileTestCase(unittest.TestCase):
         hs = yield setup_test_homeserver(
             self.addCleanup,
             http_client=None,
-            handlers=None,
             resource_for_federation=Mock(),
             federation_client=self.mock_federation,
             federation_server=Mock(),
diff --git a/tests/handlers/test_register.py b/tests/handlers/test_register.py
index cb7c0ed5..bdf3d0a8 100644
--- a/tests/handlers/test_register.py
+++ b/tests/handlers/test_register.py
@@ -18,7 +18,6 @@ from mock import Mock
 from synapse.api.auth import Auth
 from synapse.api.constants import UserTypes
 from synapse.api.errors import Codes, ResourceLimitError, SynapseError
-from synapse.handlers.register import RegistrationHandler
 from synapse.spam_checker_api import RegistrationBehaviour
 from synapse.types import RoomAlias, UserID, create_requester
 
@@ -29,11 +28,6 @@ from tests.utils import mock_getRawHeaders
 from .. import unittest
 
 
-class RegistrationHandlers:
-    def __init__(self, hs):
-        self.registration_handler = RegistrationHandler(hs)
-
-
 class RegistrationTestCase(unittest.HomeserverTestCase):
     """ Tests the RegistrationHandler. """
 
@@ -154,7 +148,7 @@ class RegistrationTestCase(unittest.HomeserverTestCase):
         room_alias_str = "#room:test"
         user_id = self.get_success(self.handler.register_user(localpart="jeff"))
         rooms = self.get_success(self.store.get_rooms_for_user(user_id))
-        directory_handler = self.hs.get_handlers().directory_handler
+        directory_handler = self.hs.get_directory_handler()
         room_alias = RoomAlias.from_string(room_alias_str)
         room_id = self.get_success(directory_handler.get_association(room_alias))
 
@@ -193,7 +187,7 @@ class RegistrationTestCase(unittest.HomeserverTestCase):
         user_id = self.get_success(self.handler.register_user(localpart="support"))
         rooms = self.get_success(self.store.get_rooms_for_user(user_id))
         self.assertEqual(len(rooms), 0)
-        directory_handler = self.hs.get_handlers().directory_handler
+        directory_handler = self.hs.get_directory_handler()
         room_alias = RoomAlias.from_string(room_alias_str)
         self.get_failure(directory_handler.get_association(room_alias), SynapseError)
 
@@ -205,7 +199,7 @@ class RegistrationTestCase(unittest.HomeserverTestCase):
         self.store.is_real_user = Mock(return_value=make_awaitable(True))
         user_id = self.get_success(self.handler.register_user(localpart="real"))
         rooms = self.get_success(self.store.get_rooms_for_user(user_id))
-        directory_handler = self.hs.get_handlers().directory_handler
+        directory_handler = self.hs.get_directory_handler()
         room_alias = RoomAlias.from_string(room_alias_str)
         room_id = self.get_success(directory_handler.get_association(room_alias))
 
@@ -237,7 +231,7 @@ class RegistrationTestCase(unittest.HomeserverTestCase):
         user_id = self.get_success(self.handler.register_user(localpart="jeff"))
 
         # Ensure the room was created.
-        directory_handler = self.hs.get_handlers().directory_handler
+        directory_handler = self.hs.get_directory_handler()
         room_alias = RoomAlias.from_string(room_alias_str)
         room_id = self.get_success(directory_handler.get_association(room_alias))
 
@@ -266,7 +260,7 @@ class RegistrationTestCase(unittest.HomeserverTestCase):
         user_id = self.get_success(self.handler.register_user(localpart="jeff"))
 
         # Ensure the room was created.
-        directory_handler = self.hs.get_handlers().directory_handler
+        directory_handler = self.hs.get_directory_handler()
         room_alias = RoomAlias.from_string(room_alias_str)
         room_id = self.get_success(directory_handler.get_association(room_alias))
 
@@ -304,7 +298,7 @@ class RegistrationTestCase(unittest.HomeserverTestCase):
         user_id = self.get_success(self.handler.register_user(localpart="jeff"))
 
         # Ensure the room was created.
-        directory_handler = self.hs.get_handlers().directory_handler
+        directory_handler = self.hs.get_directory_handler()
         room_alias = RoomAlias.from_string(room_alias_str)
         room_id = self.get_success(directory_handler.get_association(room_alias))
 
@@ -347,7 +341,7 @@ class RegistrationTestCase(unittest.HomeserverTestCase):
         )
 
         # Ensure the room was created.
-        directory_handler = self.hs.get_handlers().directory_handler
+        directory_handler = self.hs.get_directory_handler()
         room_alias = RoomAlias.from_string(room_alias_str)
         room_id = self.get_success(directory_handler.get_association(room_alias))
 
@@ -384,7 +378,7 @@ class RegistrationTestCase(unittest.HomeserverTestCase):
         user_id = self.get_success(self.handler.register_user(localpart="jeff"))
 
         # Ensure the room was created.
-        directory_handler = self.hs.get_handlers().directory_handler
+        directory_handler = self.hs.get_directory_handler()
         room_alias = RoomAlias.from_string(room_alias_str)
         room_id = self.get_success(directory_handler.get_association(room_alias))
 
@@ -413,7 +407,7 @@ class RegistrationTestCase(unittest.HomeserverTestCase):
             )
         )
         self.get_success(
-            event_creation_handler.send_nonmember_event(requester, event, context)
+            event_creation_handler.handle_new_client_event(requester, event, context)
         )
 
         # Register a second user, which won't be be in the room (or even have an invite)
diff --git a/tests/handlers/test_typing.py b/tests/handlers/test_typing.py
index 3fec09ea..16ff2e22 100644
--- a/tests/handlers/test_typing.py
+++ b/tests/handlers/test_typing.py
@@ -65,26 +65,6 @@ class TypingNotificationsTestCase(unittest.HomeserverTestCase):
         mock_federation_client = Mock(spec=["put_json"])
         mock_federation_client.put_json.return_value = defer.succeed((200, "OK"))
 
-        datastores = Mock()
-        datastores.main = Mock(
-            spec=[
-                # Bits that Federation needs
-                "prep_send_transaction",
-                "delivered_txn",
-                "get_received_txn_response",
-                "set_received_txn_response",
-                "get_destination_last_successful_stream_ordering",
-                "get_destination_retry_timings",
-                "get_devices_by_remote",
-                "maybe_store_room_on_invite",
-                # Bits that user_directory needs
-                "get_user_directory_stream_pos",
-                "get_current_state_deltas",
-                "get_device_updates_by_remote",
-                "get_room_max_stream_ordering",
-            ]
-        )
-
         # the tests assume that we are starting at unix time 1000
         reactor.pump((1000,))
 
@@ -95,8 +75,6 @@ class TypingNotificationsTestCase(unittest.HomeserverTestCase):
             replication_streams={},
         )
 
-        hs.datastores = datastores
-
         return hs
 
     def prepare(self, reactor, clock, hs):
@@ -114,16 +92,16 @@ class TypingNotificationsTestCase(unittest.HomeserverTestCase):
             "retry_interval": 0,
             "failure_ts": None,
         }
-        self.datastore.get_destination_retry_timings.return_value = defer.succeed(
-            retry_timings_res
+        self.datastore.get_destination_retry_timings = Mock(
+            return_value=defer.succeed(retry_timings_res)
         )
 
-        self.datastore.get_device_updates_by_remote.return_value = make_awaitable(
-            (0, [])
+        self.datastore.get_device_updates_by_remote = Mock(
+            return_value=make_awaitable((0, []))
         )
 
-        self.datastore.get_destination_last_successful_stream_ordering.return_value = make_awaitable(
-            None
+        self.datastore.get_destination_last_successful_stream_ordering = Mock(
+            return_value=make_awaitable(None)
         )
 
         def get_received_txn_response(*args):
@@ -145,17 +123,19 @@ class TypingNotificationsTestCase(unittest.HomeserverTestCase):
 
         self.datastore.get_joined_hosts_for_room = get_joined_hosts_for_room
 
-        def get_users_in_room(room_id):
-            return defer.succeed({str(u) for u in self.room_members})
+        async def get_users_in_room(room_id):
+            return {str(u) for u in self.room_members}
 
         self.datastore.get_users_in_room = get_users_in_room
 
-        self.datastore.get_user_directory_stream_pos.side_effect = (
-            # we deliberately return a non-None stream pos to avoid doing an initial_spam
-            lambda: make_awaitable(1)
+        self.datastore.get_user_directory_stream_pos = Mock(
+            side_effect=(
+                # we deliberately return a non-None stream pos to avoid doing an initial_spam
+                lambda: make_awaitable(1)
+            )
         )
 
-        self.datastore.get_current_state_deltas.return_value = (0, None)
+        self.datastore.get_current_state_deltas = Mock(return_value=(0, None))
 
         self.datastore.get_to_device_stream_token = lambda: 0
         self.datastore.get_new_device_msgs_for_remote = lambda *args, **kargs: make_awaitable(
diff --git a/tests/logging/test_terse_json.py b/tests/logging/test_terse_json.py
index 4cf81f71..fd128b88 100644
--- a/tests/logging/test_terse_json.py
+++ b/tests/logging/test_terse_json.py
@@ -78,7 +78,7 @@ class TerseJSONTCPTestCase(StructuredLoggingTestBase, HomeserverTestCase):
             "server_name",
             "name",
         ]
-        self.assertEqual(set(log.keys()), set(expected_log_keys))
+        self.assertCountEqual(log.keys(), expected_log_keys)
 
         # It contains the data we expect.
         self.assertEqual(log["name"], "wally")
diff --git a/tests/module_api/test_api.py b/tests/module_api/test_api.py
index 04de0b9d..9b573ac2 100644
--- a/tests/module_api/test_api.py
+++ b/tests/module_api/test_api.py
@@ -12,16 +12,27 @@
 # 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 mock import Mock
 
-from synapse.module_api import ModuleApi
+from synapse.events import EventBase
+from synapse.rest import admin
+from synapse.rest.client.v1 import login, room
+from synapse.types import create_requester
 
 from tests.unittest import HomeserverTestCase
 
 
 class ModuleApiTestCase(HomeserverTestCase):
+    servlets = [
+        admin.register_servlets,
+        login.register_servlets,
+        room.register_servlets,
+    ]
+
     def prepare(self, reactor, clock, homeserver):
         self.store = homeserver.get_datastore()
-        self.module_api = ModuleApi(homeserver, homeserver.get_auth_handler())
+        self.module_api = homeserver.get_module_api()
+        self.event_creation_handler = homeserver.get_event_creation_handler()
 
     def test_can_register_user(self):
         """Tests that an external module can register a user"""
@@ -52,3 +63,141 @@ class ModuleApiTestCase(HomeserverTestCase):
         # Check that the displayname was assigned
         displayname = self.get_success(self.store.get_profile_displayname("bob"))
         self.assertEqual(displayname, "Bobberino")
+
+    def test_sending_events_into_room(self):
+        """Tests that a module can send events into a room"""
+        # Mock out create_and_send_nonmember_event to check whether events are being sent
+        self.event_creation_handler.create_and_send_nonmember_event = Mock(
+            spec=[],
+            side_effect=self.event_creation_handler.create_and_send_nonmember_event,
+        )
+
+        # Create a user and room to play with
+        user_id = self.register_user("summer", "monkey")
+        tok = self.login("summer", "monkey")
+        room_id = self.helper.create_room_as(user_id, tok=tok)
+
+        # Create and send a non-state event
+        content = {"body": "I am a puppet", "msgtype": "m.text"}
+        event_dict = {
+            "room_id": room_id,
+            "type": "m.room.message",
+            "content": content,
+            "sender": user_id,
+        }
+        event = self.get_success(
+            self.module_api.create_and_send_event_into_room(event_dict)
+        )  # type: EventBase
+        self.assertEqual(event.sender, user_id)
+        self.assertEqual(event.type, "m.room.message")
+        self.assertEqual(event.room_id, room_id)
+        self.assertFalse(hasattr(event, "state_key"))
+        self.assertDictEqual(event.content, content)
+
+        # Check that the event was sent
+        self.event_creation_handler.create_and_send_nonmember_event.assert_called_with(
+            create_requester(user_id),
+            event_dict,
+            ratelimit=False,
+            ignore_shadow_ban=True,
+        )
+
+        # Create and send a state event
+        content = {
+            "events_default": 0,
+            "users": {user_id: 100},
+            "state_default": 50,
+            "users_default": 0,
+            "events": {"test.event.type": 25},
+        }
+        event_dict = {
+            "room_id": room_id,
+            "type": "m.room.power_levels",
+            "content": content,
+            "sender": user_id,
+            "state_key": "",
+        }
+        event = self.get_success(
+            self.module_api.create_and_send_event_into_room(event_dict)
+        )  # type: EventBase
+        self.assertEqual(event.sender, user_id)
+        self.assertEqual(event.type, "m.room.power_levels")
+        self.assertEqual(event.room_id, room_id)
+        self.assertEqual(event.state_key, "")
+        self.assertDictEqual(event.content, content)
+
+        # Check that the event was sent
+        self.event_creation_handler.create_and_send_nonmember_event.assert_called_with(
+            create_requester(user_id),
+            {
+                "type": "m.room.power_levels",
+                "content": content,
+                "room_id": room_id,
+                "sender": user_id,
+                "state_key": "",
+            },
+            ratelimit=False,
+            ignore_shadow_ban=True,
+        )
+
+        # Check that we can't send membership events
+        content = {
+            "membership": "leave",
+        }
+        event_dict = {
+            "room_id": room_id,
+            "type": "m.room.member",
+            "content": content,
+            "sender": user_id,
+            "state_key": user_id,
+        }
+        self.get_failure(
+            self.module_api.create_and_send_event_into_room(event_dict), Exception
+        )
+
+    def test_public_rooms(self):
+        """Tests that a room can be added and removed from the public rooms list,
+        as well as have its public rooms directory state queried.
+        """
+        # Create a user and room to play with
+        user_id = self.register_user("kermit", "monkey")
+        tok = self.login("kermit", "monkey")
+        room_id = self.helper.create_room_as(user_id, tok=tok)
+
+        # The room should not currently be in the public rooms directory
+        is_in_public_rooms = self.get_success(
+            self.module_api.public_room_list_manager.room_is_in_public_room_list(
+                room_id
+            )
+        )
+        self.assertFalse(is_in_public_rooms)
+
+        # Let's try adding it to the public rooms directory
+        self.get_success(
+            self.module_api.public_room_list_manager.add_room_to_public_room_list(
+                room_id
+            )
+        )
+
+        # And checking whether it's in there...
+        is_in_public_rooms = self.get_success(
+            self.module_api.public_room_list_manager.room_is_in_public_room_list(
+                room_id
+            )
+        )
+        self.assertTrue(is_in_public_rooms)
+
+        # Let's remove it again
+        self.get_success(
+            self.module_api.public_room_list_manager.remove_room_from_public_room_list(
+                room_id
+            )
+        )
+
+        # Should be gone
+        is_in_public_rooms = self.get_success(
+            self.module_api.public_room_list_manager.room_is_in_public_room_list(
+                room_id
+            )
+        )
+        self.assertFalse(is_in_public_rooms)
diff --git a/tests/push/test_email.py b/tests/push/test_email.py
index 32245686..55545d93 100644
--- a/tests/push/test_email.py
+++ b/tests/push/test_email.py
@@ -158,8 +158,21 @@ class EmailPusherTests(HomeserverTestCase):
         # We should get emailed about those messages
         self._check_for_mail()
 
+    def test_encrypted_message(self):
+        room = self.helper.create_room_as(self.user_id, tok=self.access_token)
+        self.helper.invite(
+            room=room, src=self.user_id, tok=self.access_token, targ=self.others[0].id
+        )
+        self.helper.join(room=room, user=self.others[0].id, tok=self.others[0].token)
+
+        # The other user sends some messages
+        self.helper.send_event(room, "m.room.encrypted", {}, tok=self.others[0].token)
+
+        # We should get emailed about that message
+        self._check_for_mail()
+
     def _check_for_mail(self):
-        "Check that the user receives an email notification"
+        """Check that the user receives an email notification"""
 
         # Get the stream ordering before it gets sent
         pushers = self.get_success(
diff --git a/tests/replication/_base.py b/tests/replication/_base.py
index ae60874e..093e2faa 100644
--- a/tests/replication/_base.py
+++ b/tests/replication/_base.py
@@ -12,13 +12,14 @@
 # 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.
-
 import logging
 from typing import Any, Callable, List, Optional, Tuple
 
 import attr
+import hiredis
 
 from twisted.internet.interfaces import IConsumer, IPullProducer, IReactorTime
+from twisted.internet.protocol import Protocol
 from twisted.internet.task import LoopingCall
 from twisted.web.http import HTTPChannel
 
@@ -27,7 +28,7 @@ from synapse.app.generic_worker import (
     GenericWorkerServer,
 )
 from synapse.http.server import JsonResource
-from synapse.http.site import SynapseRequest
+from synapse.http.site import SynapseRequest, SynapseSite
 from synapse.replication.http import ReplicationRestResource, streams
 from synapse.replication.tcp.handler import ReplicationCommandHandler
 from synapse.replication.tcp.protocol import ClientReplicationStreamProtocol
@@ -58,7 +59,7 @@ class BaseStreamTestCase(unittest.HomeserverTestCase):
         self.reactor.lookups["testserv"] = "1.2.3.4"
         self.worker_hs = self.setup_test_homeserver(
             http_client=None,
-            homeserverToUse=GenericWorkerServer,
+            homeserver_to_use=GenericWorkerServer,
             config=self._get_worker_hs_config(),
             reactor=self.reactor,
         )
@@ -197,19 +198,37 @@ class BaseMultiWorkerStreamTestCase(unittest.HomeserverTestCase):
         self.server_factory = ReplicationStreamProtocolFactory(self.hs)
         self.streamer = self.hs.get_replication_streamer()
 
+        # Fake in memory Redis server that servers can connect to.
+        self._redis_server = FakeRedisPubSubServer()
+
         store = self.hs.get_datastore()
         self.database_pool = store.db_pool
 
         self.reactor.lookups["testserv"] = "1.2.3.4"
+        self.reactor.lookups["localhost"] = "127.0.0.1"
+
+        # A map from a HS instance to the associated HTTP Site to use for
+        # handling inbound HTTP requests to that instance.
+        self._hs_to_site = {self.hs: self.site}
+
+        if self.hs.config.redis.redis_enabled:
+            # Handle attempts to connect to fake redis server.
+            self.reactor.add_tcp_client_callback(
+                "localhost", 6379, self.connect_any_redis_attempts,
+            )
 
-        self._worker_hs_to_resource = {}
+            self.hs.get_tcp_replication().start_replication(self.hs)
 
         # When we see a connection attempt to the master replication listener we
         # automatically set up the connection. This is so that tests don't
         # manually have to go and explicitly set it up each time (plus sometimes
         # it is impossible to write the handling explicitly in the tests).
+        #
+        # Register the master replication listener:
         self.reactor.add_tcp_client_callback(
-            "1.2.3.4", 8765, self._handle_http_replication_attempt
+            "1.2.3.4",
+            8765,
+            lambda: self._handle_http_replication_attempt(self.hs, 8765),
         )
 
     def create_test_json_resource(self):
@@ -247,34 +266,69 @@ class BaseMultiWorkerStreamTestCase(unittest.HomeserverTestCase):
         config.update(extra_config)
 
         worker_hs = self.setup_test_homeserver(
-            homeserverToUse=GenericWorkerServer,
+            homeserver_to_use=GenericWorkerServer,
             config=config,
             reactor=self.reactor,
             **kwargs
         )
 
+        # If the instance is in the `instance_map` config then workers may try
+        # and send HTTP requests to it, so we register it with
+        # `_handle_http_replication_attempt` like we do with the master HS.
+        instance_name = worker_hs.get_instance_name()
+        instance_loc = worker_hs.config.worker.instance_map.get(instance_name)
+        if instance_loc:
+            # Ensure the host is one that has a fake DNS entry.
+            if instance_loc.host not in self.reactor.lookups:
+                raise Exception(
+                    "Host does not have an IP for instance_map[%r].host = %r"
+                    % (instance_name, instance_loc.host,)
+                )
+
+            self.reactor.add_tcp_client_callback(
+                self.reactor.lookups[instance_loc.host],
+                instance_loc.port,
+                lambda: self._handle_http_replication_attempt(
+                    worker_hs, instance_loc.port
+                ),
+            )
+
         store = worker_hs.get_datastore()
         store.db_pool._db_pool = self.database_pool._db_pool
 
-        repl_handler = ReplicationCommandHandler(worker_hs)
-        client = ClientReplicationStreamProtocol(
-            worker_hs, "client", "test", self.clock, repl_handler,
-        )
-        server = self.server_factory.buildProtocol(None)
+        # Set up TCP replication between master and the new worker if we don't
+        # have Redis support enabled.
+        if not worker_hs.config.redis_enabled:
+            repl_handler = ReplicationCommandHandler(worker_hs)
+            client = ClientReplicationStreamProtocol(
+                worker_hs, "client", "test", self.clock, repl_handler,
+            )
+            server = self.server_factory.buildProtocol(None)
 
-        client_transport = FakeTransport(server, self.reactor)
-        client.makeConnection(client_transport)
+            client_transport = FakeTransport(server, self.reactor)
+            client.makeConnection(client_transport)
 
-        server_transport = FakeTransport(client, self.reactor)
-        server.makeConnection(server_transport)
+            server_transport = FakeTransport(client, self.reactor)
+            server.makeConnection(server_transport)
 
         # Set up a resource for the worker
-        resource = ReplicationRestResource(self.hs)
+        resource = ReplicationRestResource(worker_hs)
 
         for servlet in self.servlets:
             servlet(worker_hs, resource)
 
-        self._worker_hs_to_resource[worker_hs] = resource
+        self._hs_to_site[worker_hs] = SynapseSite(
+            logger_name="synapse.access.http.fake",
+            site_tag="{}-{}".format(
+                worker_hs.config.server.server_name, worker_hs.get_instance_name()
+            ),
+            config=worker_hs.config.server.listeners[0],
+            resource=resource,
+            server_version_string="1",
+        )
+
+        if worker_hs.config.redis.redis_enabled:
+            worker_hs.get_tcp_replication().start_replication(worker_hs)
 
         return worker_hs
 
@@ -285,7 +339,7 @@ class BaseMultiWorkerStreamTestCase(unittest.HomeserverTestCase):
         return config
 
     def render_on_worker(self, worker_hs: HomeServer, request: SynapseRequest):
-        render(request, self._worker_hs_to_resource[worker_hs], self.reactor)
+        render(request, self._hs_to_site[worker_hs].resource, self.reactor)
 
     def replicate(self):
         """Tell the master side of replication that something has happened, and then
@@ -294,9 +348,9 @@ class BaseMultiWorkerStreamTestCase(unittest.HomeserverTestCase):
         self.streamer.on_notifier_poke()
         self.pump()
 
-    def _handle_http_replication_attempt(self):
-        """Handles a connection attempt to the master replication HTTP
-        listener.
+    def _handle_http_replication_attempt(self, hs, repl_port):
+        """Handles a connection attempt to the given HS replication HTTP
+        listener on the given port.
         """
 
         # We should have at least one outbound connection attempt, where the
@@ -305,7 +359,7 @@ class BaseMultiWorkerStreamTestCase(unittest.HomeserverTestCase):
         self.assertGreaterEqual(len(clients), 1)
         (host, port, client_factory, _timeout, _bindAddress) = clients.pop()
         self.assertEqual(host, "1.2.3.4")
-        self.assertEqual(port, 8765)
+        self.assertEqual(port, repl_port)
 
         # Set up client side protocol
         client_protocol = client_factory.buildProtocol(None)
@@ -315,7 +369,7 @@ class BaseMultiWorkerStreamTestCase(unittest.HomeserverTestCase):
         # Set up the server side protocol
         channel = _PushHTTPChannel(self.reactor)
         channel.requestFactory = request_factory
-        channel.site = self.site
+        channel.site = self._hs_to_site[hs]
 
         # Connect client to server and vice versa.
         client_to_server_transport = FakeTransport(
@@ -333,6 +387,32 @@ class BaseMultiWorkerStreamTestCase(unittest.HomeserverTestCase):
         # inside `connecTCP` before the connection has been passed back to the
         # code that requested the TCP connection.
 
+    def connect_any_redis_attempts(self):
+        """If redis is enabled we need to deal with workers connecting to a
+        redis server. We don't want to use a real Redis server so we use a
+        fake one.
+        """
+        clients = self.reactor.tcpClients
+        self.assertEqual(len(clients), 1)
+        (host, port, client_factory, _timeout, _bindAddress) = clients.pop(0)
+        self.assertEqual(host, "localhost")
+        self.assertEqual(port, 6379)
+
+        client_protocol = client_factory.buildProtocol(None)
+        server_protocol = self._redis_server.buildProtocol(None)
+
+        client_to_server_transport = FakeTransport(
+            server_protocol, self.reactor, client_protocol
+        )
+        client_protocol.makeConnection(client_to_server_transport)
+
+        server_to_client_transport = FakeTransport(
+            client_protocol, self.reactor, server_protocol
+        )
+        server_protocol.makeConnection(server_to_client_transport)
+
+        return client_to_server_transport, server_to_client_transport
+
 
 class TestReplicationDataHandler(GenericWorkerReplicationHandler):
     """Drop-in for ReplicationDataHandler which just collects RDATA rows"""
@@ -467,3 +547,105 @@ class _PullToPushProducer:
                 pass
 
             self.stopProducing()
+
+
+class FakeRedisPubSubServer:
+    """A fake Redis server for pub/sub.
+    """
+
+    def __init__(self):
+        self._subscribers = set()
+
+    def add_subscriber(self, conn):
+        """A connection has called SUBSCRIBE
+        """
+        self._subscribers.add(conn)
+
+    def remove_subscriber(self, conn):
+        """A connection has called UNSUBSCRIBE
+        """
+        self._subscribers.discard(conn)
+
+    def publish(self, conn, channel, msg) -> int:
+        """A connection want to publish a message to subscribers.
+        """
+        for sub in self._subscribers:
+            sub.send(["message", channel, msg])
+
+        return len(self._subscribers)
+
+    def buildProtocol(self, addr):
+        return FakeRedisPubSubProtocol(self)
+
+
+class FakeRedisPubSubProtocol(Protocol):
+    """A connection from a client talking to the fake Redis server.
+    """
+
+    def __init__(self, server: FakeRedisPubSubServer):
+        self._server = server
+        self._reader = hiredis.Reader()
+
+    def dataReceived(self, data):
+        self._reader.feed(data)
+
+        # We might get multiple messages in one packet.
+        while True:
+            msg = self._reader.gets()
+
+            if msg is False:
+                # No more messages.
+                return
+
+            if not isinstance(msg, list):
+                # Inbound commands should always be a list
+                raise Exception("Expected redis list")
+
+            self.handle_command(msg[0], *msg[1:])
+
+    def handle_command(self, command, *args):
+        """Received a Redis command from the client.
+        """
+
+        # We currently only support pub/sub.
+        if command == b"PUBLISH":
+            channel, message = args
+            num_subscribers = self._server.publish(self, channel, message)
+            self.send(num_subscribers)
+        elif command == b"SUBSCRIBE":
+            (channel,) = args
+            self._server.add_subscriber(self)
+            self.send(["subscribe", channel, 1])
+        else:
+            raise Exception("Unknown command")
+
+    def send(self, msg):
+        """Send a message back to the client.
+        """
+        raw = self.encode(msg).encode("utf-8")
+
+        self.transport.write(raw)
+        self.transport.flush()
+
+    def encode(self, obj):
+        """Encode an object to its Redis format.
+
+        Supports: strings/bytes, integers and list/tuples.
+        """
+
+        if isinstance(obj, bytes):
+            # We assume bytes are just unicode strings.
+            obj = obj.decode("utf-8")
+
+        if isinstance(obj, str):
+            return "${len}\r\n{str}\r\n".format(len=len(obj), str=obj)
+        if isinstance(obj, int):
+            return ":{val}\r\n".format(val=obj)
+        if isinstance(obj, (list, tuple)):
+            items = "".join(self.encode(a) for a in obj)
+            return "*{len}\r\n{items}".format(len=len(obj), items=items)
+
+        raise Exception("Unrecognized type for encoding redis: %r: %r", type(obj), obj)
+
+    def connectionLost(self, reason):
+        self._server.remove_subscriber(self)
diff --git a/tests/replication/test_federation_ack.py b/tests/replication/test_federation_ack.py
index 23be1167..18536675 100644
--- a/tests/replication/test_federation_ack.py
+++ b/tests/replication/test_federation_ack.py
@@ -31,7 +31,7 @@ class FederationAckTestCase(HomeserverTestCase):
         return config
 
     def make_homeserver(self, reactor, clock):
-        hs = self.setup_test_homeserver(homeserverToUse=GenericWorkerServer)
+        hs = self.setup_test_homeserver(homeserver_to_use=GenericWorkerServer)
 
         return hs
 
diff --git a/tests/replication/test_federation_sender_shard.py b/tests/replication/test_federation_sender_shard.py
index 1d7edee5..779745ae 100644
--- a/tests/replication/test_federation_sender_shard.py
+++ b/tests/replication/test_federation_sender_shard.py
@@ -207,7 +207,7 @@ class FederationSenderTestCase(BaseMultiWorkerStreamTestCase):
     def create_room_with_remote_server(self, user, token, remote_server="other_server"):
         room = self.helper.create_room_as(user, tok=token)
         store = self.hs.get_datastore()
-        federation = self.hs.get_handlers().federation_handler
+        federation = self.hs.get_federation_handler()
 
         prev_event_ids = self.get_success(store.get_latest_event_ids_in_room(room))
         room_version = self.get_success(store.get_room_version(room))
@@ -226,7 +226,7 @@ class FederationSenderTestCase(BaseMultiWorkerStreamTestCase):
         }
 
         builder = factory.for_room_version(room_version, event_dict)
-        join_event = self.get_success(builder.build(prev_event_ids))
+        join_event = self.get_success(builder.build(prev_event_ids, None))
 
         self.get_success(federation.on_send_join_request(remote_server, join_event))
         self.replicate()
diff --git a/tests/replication/test_sharded_event_persister.py b/tests/replication/test_sharded_event_persister.py
new file mode 100644
index 00000000..82cf033d
--- /dev/null
+++ b/tests/replication/test_sharded_event_persister.py
@@ -0,0 +1,319 @@
+# -*- coding: utf-8 -*-
+# Copyright 2020 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.
+import logging
+
+from mock import patch
+
+from synapse.api.room_versions import RoomVersion
+from synapse.rest import admin
+from synapse.rest.client.v1 import login, room
+from synapse.rest.client.v2_alpha import sync
+
+from tests.replication._base import BaseMultiWorkerStreamTestCase
+from tests.utils import USE_POSTGRES_FOR_TESTS
+
+logger = logging.getLogger(__name__)
+
+
+class EventPersisterShardTestCase(BaseMultiWorkerStreamTestCase):
+    """Checks event persisting sharding works
+    """
+
+    # Event persister sharding requires postgres (due to needing
+    # `MutliWriterIdGenerator`).
+    if not USE_POSTGRES_FOR_TESTS:
+        skip = "Requires Postgres"
+
+    servlets = [
+        admin.register_servlets_for_client_rest_resource,
+        room.register_servlets,
+        login.register_servlets,
+        sync.register_servlets,
+    ]
+
+    def prepare(self, reactor, clock, hs):
+        # Register a user who sends a message that we'll get notified about
+        self.other_user_id = self.register_user("otheruser", "pass")
+        self.other_access_token = self.login("otheruser", "pass")
+
+        self.room_creator = self.hs.get_room_creation_handler()
+        self.store = hs.get_datastore()
+
+    def default_config(self):
+        conf = super().default_config()
+        conf["redis"] = {"enabled": "true"}
+        conf["stream_writers"] = {"events": ["worker1", "worker2"]}
+        conf["instance_map"] = {
+            "worker1": {"host": "testserv", "port": 1001},
+            "worker2": {"host": "testserv", "port": 1002},
+        }
+        return conf
+
+    def _create_room(self, room_id: str, user_id: str, tok: str):
+        """Create a room with given room_id
+        """
+
+        # We control the room ID generation by patching out the
+        # `_generate_room_id` method
+        async def generate_room(
+            creator_id: str, is_public: bool, room_version: RoomVersion
+        ):
+            await self.store.store_room(
+                room_id=room_id,
+                room_creator_user_id=creator_id,
+                is_public=is_public,
+                room_version=room_version,
+            )
+            return room_id
+
+        with patch(
+            "synapse.handlers.room.RoomCreationHandler._generate_room_id"
+        ) as mock:
+            mock.side_effect = generate_room
+            self.helper.create_room_as(user_id, tok=tok)
+
+    def test_basic(self):
+        """Simple test to ensure that multiple rooms can be created and joined,
+        and that different rooms get handled by different instances.
+        """
+
+        self.make_worker_hs(
+            "synapse.app.generic_worker", {"worker_name": "worker1"},
+        )
+
+        self.make_worker_hs(
+            "synapse.app.generic_worker", {"worker_name": "worker2"},
+        )
+
+        persisted_on_1 = False
+        persisted_on_2 = False
+
+        store = self.hs.get_datastore()
+
+        user_id = self.register_user("user", "pass")
+        access_token = self.login("user", "pass")
+
+        # Keep making new rooms until we see rooms being persisted on both
+        # workers.
+        for _ in range(10):
+            # Create a room
+            room = self.helper.create_room_as(user_id, tok=access_token)
+
+            # The other user joins
+            self.helper.join(
+                room=room, user=self.other_user_id, tok=self.other_access_token
+            )
+
+            # The other user sends some messages
+            rseponse = self.helper.send(room, body="Hi!", tok=self.other_access_token)
+            event_id = rseponse["event_id"]
+
+            # The event position includes which instance persisted the event.
+            pos = self.get_success(store.get_position_for_event(event_id))
+
+            persisted_on_1 |= pos.instance_name == "worker1"
+            persisted_on_2 |= pos.instance_name == "worker2"
+
+            if persisted_on_1 and persisted_on_2:
+                break
+
+        self.assertTrue(persisted_on_1)
+        self.assertTrue(persisted_on_2)
+
+    def test_vector_clock_token(self):
+        """Tests that using a stream token with a vector clock component works
+        correctly with basic /sync and /messages usage.
+        """
+
+        self.make_worker_hs(
+            "synapse.app.generic_worker", {"worker_name": "worker1"},
+        )
+
+        worker_hs2 = self.make_worker_hs(
+            "synapse.app.generic_worker", {"worker_name": "worker2"},
+        )
+
+        sync_hs = self.make_worker_hs(
+            "synapse.app.generic_worker", {"worker_name": "sync"},
+        )
+
+        # Specially selected room IDs that get persisted on different workers.
+        room_id1 = "!foo:test"
+        room_id2 = "!baz:test"
+
+        self.assertEqual(
+            self.hs.config.worker.events_shard_config.get_instance(room_id1), "worker1"
+        )
+        self.assertEqual(
+            self.hs.config.worker.events_shard_config.get_instance(room_id2), "worker2"
+        )
+
+        user_id = self.register_user("user", "pass")
+        access_token = self.login("user", "pass")
+
+        store = self.hs.get_datastore()
+
+        # Create two room on the different workers.
+        self._create_room(room_id1, user_id, access_token)
+        self._create_room(room_id2, user_id, access_token)
+
+        # The other user joins
+        self.helper.join(
+            room=room_id1, user=self.other_user_id, tok=self.other_access_token
+        )
+        self.helper.join(
+            room=room_id2, user=self.other_user_id, tok=self.other_access_token
+        )
+
+        # Do an initial sync so that we're up to date.
+        request, channel = self.make_request("GET", "/sync", access_token=access_token)
+        self.render_on_worker(sync_hs, request)
+        next_batch = channel.json_body["next_batch"]
+
+        # We now gut wrench into the events stream MultiWriterIdGenerator on
+        # worker2 to mimic it getting stuck persisting an event. This ensures
+        # that when we send an event on worker1 we end up in a state where
+        # worker2 events stream position lags that on worker1, resulting in a
+        # RoomStreamToken with a non-empty instance map component.
+        #
+        # Worker2's event stream position will not advance until we call
+        # __aexit__ again.
+        actx = worker_hs2.get_datastore()._stream_id_gen.get_next()
+        self.get_success(actx.__aenter__())
+
+        response = self.helper.send(room_id1, body="Hi!", tok=self.other_access_token)
+        first_event_in_room1 = response["event_id"]
+
+        # Assert that the current stream token has an instance map component, as
+        # we are trying to test vector clock tokens.
+        room_stream_token = store.get_room_max_token()
+        self.assertNotEqual(len(room_stream_token.instance_map), 0)
+
+        # Check that syncing still gets the new event, despite the gap in the
+        # stream IDs.
+        request, channel = self.make_request(
+            "GET", "/sync?since={}".format(next_batch), access_token=access_token
+        )
+        self.render_on_worker(sync_hs, request)
+
+        # We should only see the new event and nothing else
+        self.assertIn(room_id1, channel.json_body["rooms"]["join"])
+        self.assertNotIn(room_id2, channel.json_body["rooms"]["join"])
+
+        events = channel.json_body["rooms"]["join"][room_id1]["timeline"]["events"]
+        self.assertListEqual(
+            [first_event_in_room1], [event["event_id"] for event in events]
+        )
+
+        # Get the next batch and makes sure its a vector clock style token.
+        vector_clock_token = channel.json_body["next_batch"]
+        self.assertTrue(vector_clock_token.startswith("m"))
+
+        # Now that we've got a vector clock token we finish the fake persisting
+        # an event we started above.
+        self.get_success(actx.__aexit__(None, None, None))
+
+        # Now try and send an event to the other rooom so that we can test that
+        # the vector clock style token works as a `since` token.
+        response = self.helper.send(room_id2, body="Hi!", tok=self.other_access_token)
+        first_event_in_room2 = response["event_id"]
+
+        request, channel = self.make_request(
+            "GET",
+            "/sync?since={}".format(vector_clock_token),
+            access_token=access_token,
+        )
+        self.render_on_worker(sync_hs, request)
+
+        self.assertNotIn(room_id1, channel.json_body["rooms"]["join"])
+        self.assertIn(room_id2, channel.json_body["rooms"]["join"])
+
+        events = channel.json_body["rooms"]["join"][room_id2]["timeline"]["events"]
+        self.assertListEqual(
+            [first_event_in_room2], [event["event_id"] for event in events]
+        )
+
+        next_batch = channel.json_body["next_batch"]
+
+        # We also want to test that the vector clock style token works with
+        # pagination. We do this by sending a couple of new events into the room
+        # and syncing again to get a prev_batch token for each room, then
+        # paginating from there back to the vector clock token.
+        self.helper.send(room_id1, body="Hi again!", tok=self.other_access_token)
+        self.helper.send(room_id2, body="Hi again!", tok=self.other_access_token)
+
+        request, channel = self.make_request(
+            "GET", "/sync?since={}".format(next_batch), access_token=access_token
+        )
+        self.render_on_worker(sync_hs, request)
+
+        prev_batch1 = channel.json_body["rooms"]["join"][room_id1]["timeline"][
+            "prev_batch"
+        ]
+        prev_batch2 = channel.json_body["rooms"]["join"][room_id2]["timeline"][
+            "prev_batch"
+        ]
+
+        # Paginating back in the first room should not produce any results, as
+        # no events have happened in it. This tests that we are correctly
+        # filtering results based on the vector clock portion.
+        request, channel = self.make_request(
+            "GET",
+            "/rooms/{}/messages?from={}&to={}&dir=b".format(
+                room_id1, prev_batch1, vector_clock_token
+            ),
+            access_token=access_token,
+        )
+        self.render_on_worker(sync_hs, request)
+        self.assertListEqual([], channel.json_body["chunk"])
+
+        # Paginating back on the second room should produce the first event
+        # again. This tests that pagination isn't completely broken.
+        request, channel = self.make_request(
+            "GET",
+            "/rooms/{}/messages?from={}&to={}&dir=b".format(
+                room_id2, prev_batch2, vector_clock_token
+            ),
+            access_token=access_token,
+        )
+        self.render_on_worker(sync_hs, request)
+        self.assertEqual(len(channel.json_body["chunk"]), 1)
+        self.assertEqual(
+            channel.json_body["chunk"][0]["event_id"], first_event_in_room2
+        )
+
+        # Paginating forwards should give the same results
+        request, channel = self.make_request(
+            "GET",
+            "/rooms/{}/messages?from={}&to={}&dir=f".format(
+                room_id1, vector_clock_token, prev_batch1
+            ),
+            access_token=access_token,
+        )
+        self.render_on_worker(sync_hs, request)
+        self.assertListEqual([], channel.json_body["chunk"])
+
+        request, channel = self.make_request(
+            "GET",
+            "/rooms/{}/messages?from={}&to={}&dir=f".format(
+                room_id2, vector_clock_token, prev_batch2,
+            ),
+            access_token=access_token,
+        )
+        self.render_on_worker(sync_hs, request)
+        self.assertEqual(len(channel.json_body["chunk"]), 1)
+        self.assertEqual(
+            channel.json_body["chunk"][0]["event_id"], first_event_in_room2
+        )
diff --git a/tests/rest/client/test_shadow_banned.py b/tests/rest/client/test_shadow_banned.py
index dfe4bf77..6bb02b96 100644
--- a/tests/rest/client/test_shadow_banned.py
+++ b/tests/rest/client/test_shadow_banned.py
@@ -78,7 +78,7 @@ class RoomTestCase(_ShadowBannedBase):
 
     def test_invite_3pid(self):
         """Ensure that a 3PID invite does not attempt to contact the identity server."""
-        identity_handler = self.hs.get_handlers().identity_handler
+        identity_handler = self.hs.get_identity_handler()
         identity_handler.lookup_3pid = Mock(
             side_effect=AssertionError("This should not get called")
         )
diff --git a/tests/rest/client/test_third_party_rules.py b/tests/rest/client/test_third_party_rules.py
new file mode 100644
index 00000000..0048bea5
--- /dev/null
+++ b/tests/rest/client/test_third_party_rules.py
@@ -0,0 +1,190 @@
+# -*- coding: utf-8 -*-
+# Copyright 2019 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.
+import threading
+from typing import Dict
+
+from mock import Mock
+
+from synapse.events import EventBase
+from synapse.module_api import ModuleApi
+from synapse.rest import admin
+from synapse.rest.client.v1 import login, room
+from synapse.types import Requester, StateMap
+
+from tests import unittest
+
+thread_local = threading.local()
+
+
+class ThirdPartyRulesTestModule:
+    def __init__(self, config: Dict, module_api: ModuleApi):
+        # keep a record of the "current" rules module, so that the test can patch
+        # it if desired.
+        thread_local.rules_module = self
+        self.module_api = module_api
+
+    async def on_create_room(
+        self, requester: Requester, config: dict, is_requester_admin: bool
+    ):
+        return True
+
+    async def check_event_allowed(self, event: EventBase, state: StateMap[EventBase]):
+        return True
+
+    @staticmethod
+    def parse_config(config):
+        return config
+
+
+def current_rules_module() -> ThirdPartyRulesTestModule:
+    return thread_local.rules_module
+
+
+class ThirdPartyRulesTestCase(unittest.HomeserverTestCase):
+    servlets = [
+        admin.register_servlets,
+        login.register_servlets,
+        room.register_servlets,
+    ]
+
+    def default_config(self):
+        config = super().default_config()
+        config["third_party_event_rules"] = {
+            "module": __name__ + ".ThirdPartyRulesTestModule",
+            "config": {},
+        }
+        return config
+
+    def prepare(self, reactor, clock, homeserver):
+        # Create a user and room to play with during the tests
+        self.user_id = self.register_user("kermit", "monkey")
+        self.tok = self.login("kermit", "monkey")
+
+        self.room_id = self.helper.create_room_as(self.user_id, tok=self.tok)
+
+    def test_third_party_rules(self):
+        """Tests that a forbidden event is forbidden from being sent, but an allowed one
+        can be sent.
+        """
+        # patch the rules module with a Mock which will return False for some event
+        # types
+        async def check(ev, state):
+            return ev.type != "foo.bar.forbidden"
+
+        callback = Mock(spec=[], side_effect=check)
+        current_rules_module().check_event_allowed = callback
+
+        request, channel = self.make_request(
+            "PUT",
+            "/_matrix/client/r0/rooms/%s/send/foo.bar.allowed/1" % self.room_id,
+            {},
+            access_token=self.tok,
+        )
+        self.render(request)
+        self.assertEquals(channel.result["code"], b"200", channel.result)
+
+        callback.assert_called_once()
+
+        # there should be various state events in the state arg: do some basic checks
+        state_arg = callback.call_args[0][1]
+        for k in (("m.room.create", ""), ("m.room.member", self.user_id)):
+            self.assertIn(k, state_arg)
+            ev = state_arg[k]
+            self.assertEqual(ev.type, k[0])
+            self.assertEqual(ev.state_key, k[1])
+
+        request, channel = self.make_request(
+            "PUT",
+            "/_matrix/client/r0/rooms/%s/send/foo.bar.forbidden/2" % self.room_id,
+            {},
+            access_token=self.tok,
+        )
+        self.render(request)
+        self.assertEquals(channel.result["code"], b"403", channel.result)
+
+    def test_cannot_modify_event(self):
+        """cannot accidentally modify an event before it is persisted"""
+
+        # first patch the event checker so that it will try to modify the event
+        async def check(ev: EventBase, state):
+            ev.content = {"x": "y"}
+            return True
+
+        current_rules_module().check_event_allowed = check
+
+        # now send the event
+        request, channel = self.make_request(
+            "PUT",
+            "/_matrix/client/r0/rooms/%s/send/modifyme/1" % self.room_id,
+            {"x": "x"},
+            access_token=self.tok,
+        )
+        self.render(request)
+        self.assertEqual(channel.result["code"], b"500", channel.result)
+
+    def test_modify_event(self):
+        """The module can return a modified version of the event"""
+        # first patch the event checker so that it will modify the event
+        async def check(ev: EventBase, state):
+            d = ev.get_dict()
+            d["content"] = {"x": "y"}
+            return d
+
+        current_rules_module().check_event_allowed = check
+
+        # now send the event
+        request, channel = self.make_request(
+            "PUT",
+            "/_matrix/client/r0/rooms/%s/send/modifyme/1" % self.room_id,
+            {"x": "x"},
+            access_token=self.tok,
+        )
+        self.render(request)
+        self.assertEqual(channel.result["code"], b"200", channel.result)
+        event_id = channel.json_body["event_id"]
+
+        # ... and check that it got modified
+        request, channel = self.make_request(
+            "GET",
+            "/_matrix/client/r0/rooms/%s/event/%s" % (self.room_id, event_id),
+            access_token=self.tok,
+        )
+        self.render(request)
+        self.assertEqual(channel.result["code"], b"200", channel.result)
+        ev = channel.json_body
+        self.assertEqual(ev["content"]["x"], "y")
+
+    def test_send_event(self):
+        """Tests that the module can send an event into a room via the module api"""
+        content = {
+            "msgtype": "m.text",
+            "body": "Hello!",
+        }
+        event_dict = {
+            "room_id": self.room_id,
+            "type": "m.room.message",
+            "content": content,
+            "sender": self.user_id,
+        }
+        event = self.get_success(
+            current_rules_module().module_api.create_and_send_event_into_room(
+                event_dict
+            )
+        )  # type: EventBase
+
+        self.assertEquals(event.sender, self.user_id)
+        self.assertEquals(event.room_id, self.room_id)
+        self.assertEquals(event.type, "m.room.message")
+        self.assertEquals(event.content, content)
diff --git a/tests/rest/client/third_party_rules.py b/tests/rest/client/third_party_rules.py
deleted file mode 100644
index 8c24add5..00000000
--- a/tests/rest/client/third_party_rules.py
+++ /dev/null
@@ -1,79 +0,0 @@
-# -*- coding: utf-8 -*-
-# Copyright 2019 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 synapse.rest import admin
-from synapse.rest.client.v1 import login, room
-
-from tests import unittest
-
-
-class ThirdPartyRulesTestModule:
-    def __init__(self, config):
-        pass
-
-    def check_event_allowed(self, event, context):
-        if event.type == "foo.bar.forbidden":
-            return False
-        else:
-            return True
-
-    @staticmethod
-    def parse_config(config):
-        return config
-
-
-class ThirdPartyRulesTestCase(unittest.HomeserverTestCase):
-    servlets = [
-        admin.register_servlets,
-        login.register_servlets,
-        room.register_servlets,
-    ]
-
-    def make_homeserver(self, reactor, clock):
-        config = self.default_config()
-        config["third_party_event_rules"] = {
-            "module": "tests.rest.client.third_party_rules.ThirdPartyRulesTestModule",
-            "config": {},
-        }
-
-        self.hs = self.setup_test_homeserver(config=config)
-        return self.hs
-
-    def test_third_party_rules(self):
-        """Tests that a forbidden event is forbidden from being sent, but an allowed one
-        can be sent.
-        """
-        user_id = self.register_user("kermit", "monkey")
-        tok = self.login("kermit", "monkey")
-
-        room_id = self.helper.create_room_as(user_id, tok=tok)
-
-        request, channel = self.make_request(
-            "PUT",
-            "/_matrix/client/r0/rooms/%s/send/foo.bar.allowed/1" % room_id,
-            {},
-            access_token=tok,
-        )
-        self.render(request)
-        self.assertEquals(channel.result["code"], b"200", channel.result)
-
-        request, channel = self.make_request(
-            "PUT",
-            "/_matrix/client/r0/rooms/%s/send/foo.bar.forbidden/1" % room_id,
-            {},
-            access_token=tok,
-        )
-        self.render(request)
-        self.assertEquals(channel.result["code"], b"403", channel.result)
diff --git a/tests/rest/client/v1/test_directory.py b/tests/rest/client/v1/test_directory.py
index 633b7dbd..ea5a7f37 100644
--- a/tests/rest/client/v1/test_directory.py
+++ b/tests/rest/client/v1/test_directory.py
@@ -21,6 +21,7 @@ from synapse.types import RoomAlias
 from synapse.util.stringutils import random_string
 
 from tests import unittest
+from tests.unittest import override_config
 
 
 class DirectoryTestCase(unittest.HomeserverTestCase):
@@ -67,10 +68,18 @@ class DirectoryTestCase(unittest.HomeserverTestCase):
         self.ensure_user_joined_room()
         self.set_alias_via_directory(400, alias_length=256)
 
-    def test_state_event_in_room(self):
+    @override_config({"default_room_version": 5})
+    def test_state_event_user_in_v5_room(self):
+        """Test that a regular user can add alias events before room v6"""
         self.ensure_user_joined_room()
         self.set_alias_via_state_event(200)
 
+    @override_config({"default_room_version": 6})
+    def test_state_event_v6_room(self):
+        """Test that a regular user can *not* add alias events from room v6"""
+        self.ensure_user_joined_room()
+        self.set_alias_via_state_event(403)
+
     def test_directory_in_room(self):
         self.ensure_user_joined_room()
         self.set_alias_via_directory(200)
diff --git a/tests/rest/client/v1/test_events.py b/tests/rest/client/v1/test_events.py
index f7552087..3397ba55 100644
--- a/tests/rest/client/v1/test_events.py
+++ b/tests/rest/client/v1/test_events.py
@@ -42,7 +42,7 @@ class EventStreamPermissionsTestCase(unittest.HomeserverTestCase):
 
         hs = self.setup_test_homeserver(config=config)
 
-        hs.get_handlers().federation_handler = Mock()
+        hs.get_federation_handler = Mock()
 
         return hs
 
diff --git a/tests/rest/client/v1/test_rooms.py b/tests/rest/client/v1/test_rooms.py
index 0d809d25..9ba5f9d9 100644
--- a/tests/rest/client/v1/test_rooms.py
+++ b/tests/rest/client/v1/test_rooms.py
@@ -32,6 +32,7 @@ from synapse.types import JsonDict, RoomAlias, UserID
 from synapse.util.stringutils import random_string
 
 from tests import unittest
+from tests.test_utils import make_awaitable
 
 PATH_PREFIX = b"/_matrix/client/api/v1"
 
@@ -47,7 +48,10 @@ class RoomBase(unittest.HomeserverTestCase):
             "red", http_client=None, federation_client=Mock(),
         )
 
-        self.hs.get_federation_handler = Mock(return_value=Mock())
+        self.hs.get_federation_handler = Mock()
+        self.hs.get_federation_handler.return_value.maybe_backfill = Mock(
+            return_value=make_awaitable(None)
+        )
 
         async def _insert_client_ip(*args, **kwargs):
             return None
diff --git a/tests/rest/client/v1/test_typing.py b/tests/rest/client/v1/test_typing.py
index 94d2bf2e..cd58ee77 100644
--- a/tests/rest/client/v1/test_typing.py
+++ b/tests/rest/client/v1/test_typing.py
@@ -44,7 +44,7 @@ class RoomTypingTestCase(unittest.HomeserverTestCase):
 
         self.event_source = hs.get_event_sources().sources["typing"]
 
-        hs.get_handlers().federation_handler = Mock()
+        hs.get_federation_handler = Mock()
 
         async def get_user_by_access_token(token=None, allow_guest=False):
             return {
diff --git a/tests/rest/client/v2_alpha/test_account.py b/tests/rest/client/v2_alpha/test_account.py
index ae2cd67f..66ac4dbe 100644
--- a/tests/rest/client/v2_alpha/test_account.py
+++ b/tests/rest/client/v2_alpha/test_account.py
@@ -352,7 +352,6 @@ class DeactivateTestCase(unittest.HomeserverTestCase):
         self.render(request)
         self.assertEqual(request.code, 401)
 
-    @unittest.INFO
     def test_pending_invites(self):
         """Tests that deactivating a user rejects every pending invite for them."""
         store = self.hs.get_datastore()
diff --git a/tests/rest/client/v2_alpha/test_auth.py b/tests/rest/client/v2_alpha/test_auth.py
index 293ccfba..86184f0d 100644
--- a/tests/rest/client/v2_alpha/test_auth.py
+++ b/tests/rest/client/v2_alpha/test_auth.py
@@ -104,7 +104,6 @@ class FallbackAuthTests(unittest.HomeserverTestCase):
         self.assertEqual(len(attempts), 1)
         self.assertEqual(attempts[0][0]["response"], "a")
 
-    @unittest.INFO
     def test_fallback_captcha(self):
         """Ensure that fallback auth via a captcha works."""
         # Returns a 401 as per the spec
diff --git a/tests/server.py b/tests/server.py
index b404ad4e..4d33b840 100644
--- a/tests/server.py
+++ b/tests/server.py
@@ -1,8 +1,11 @@
 import json
 import logging
+from collections import deque
 from io import SEEK_END, BytesIO
+from typing import Callable
 
 import attr
+from typing_extensions import Deque
 from zope.interface import implementer
 
 from twisted.internet import address, threads, udp
@@ -251,6 +254,7 @@ class ThreadedMemoryReactorClock(MemoryReactorClock):
         self._tcp_callbacks = {}
         self._udp = []
         lookups = self.lookups = {}
+        self._thread_callbacks = deque()  # type: Deque[Callable[[], None]]()
 
         @implementer(IResolverSimple)
         class FakeResolver:
@@ -272,10 +276,10 @@ class ThreadedMemoryReactorClock(MemoryReactorClock):
         """
         Make the callback fire in the next reactor iteration.
         """
-        d = Deferred()
-        d.addCallback(lambda x: callback(*args, **kwargs))
-        self.callLater(0, d.callback, True)
-        return d
+        cb = lambda: callback(*args, **kwargs)
+        # it's not safe to call callLater() here, so we append the callback to a
+        # separate queue.
+        self._thread_callbacks.append(cb)
 
     def getThreadPool(self):
         return self.threadpool
@@ -303,6 +307,30 @@ class ThreadedMemoryReactorClock(MemoryReactorClock):
 
         return conn
 
+    def advance(self, amount):
+        # first advance our reactor's time, and run any "callLater" callbacks that
+        # makes ready
+        super().advance(amount)
+
+        # now run any "callFromThread" callbacks
+        while True:
+            try:
+                callback = self._thread_callbacks.popleft()
+            except IndexError:
+                break
+            callback()
+
+            # check for more "callLater" callbacks added by the thread callback
+            # This isn't required in a regular reactor, but it ends up meaning that
+            # our database queries can complete in a single call to `advance` [1] which
+            # simplifies tests.
+            #
+            # [1]: we replace the threadpool backing the db connection pool with a
+            # mock ThreadPool which doesn't really use threads; but we still use
+            # reactor.callFromThread to feed results back from the db functions to the
+            # main thread.
+            super().advance(0)
+
 
 class ThreadPool:
     """
@@ -339,8 +367,6 @@ def setup_test_homeserver(cleanup_func, *args, **kwargs):
     """
     server = _sth(cleanup_func, *args, **kwargs)
 
-    database = server.config.database.get_single_database()
-
     # Make the thread pool synchronous.
     clock = server.get_clock()
 
@@ -372,6 +398,10 @@ def setup_test_homeserver(cleanup_func, *args, **kwargs):
         pool.threadpool = ThreadPool(clock._reactor)
         pool.running = True
 
+    # We've just changed the Databases to run DB transactions on the same
+    # thread, so we need to disable the dedicated thread behaviour.
+    server.get_datastores().main.USE_DEDICATED_DB_THREADS_FOR_EVENT_FETCHING = False
+
     return server
 
 
diff --git a/tests/storage/test__base.py b/tests/storage/test__base.py
index f5afed01..1ac4ebc6 100644
--- a/tests/storage/test__base.py
+++ b/tests/storage/test__base.py
@@ -15,308 +15,9 @@
 # limitations under the License.
 
 
-from mock import Mock
-
-from twisted.internet import defer
-
-from synapse.util.async_helpers import ObservableDeferred
-from synapse.util.caches.descriptors import Cache, cached
-
 from tests import unittest
 
 
-class CacheTestCase(unittest.HomeserverTestCase):
-    def prepare(self, reactor, clock, homeserver):
-        self.cache = Cache("test")
-
-    def test_empty(self):
-        failed = False
-        try:
-            self.cache.get("foo")
-        except KeyError:
-            failed = True
-
-        self.assertTrue(failed)
-
-    def test_hit(self):
-        self.cache.prefill("foo", 123)
-
-        self.assertEquals(self.cache.get("foo"), 123)
-
-    def test_invalidate(self):
-        self.cache.prefill(("foo",), 123)
-        self.cache.invalidate(("foo",))
-
-        failed = False
-        try:
-            self.cache.get(("foo",))
-        except KeyError:
-            failed = True
-
-        self.assertTrue(failed)
-
-    def test_eviction(self):
-        cache = Cache("test", max_entries=2)
-
-        cache.prefill(1, "one")
-        cache.prefill(2, "two")
-        cache.prefill(3, "three")  # 1 will be evicted
-
-        failed = False
-        try:
-            cache.get(1)
-        except KeyError:
-            failed = True
-
-        self.assertTrue(failed)
-
-        cache.get(2)
-        cache.get(3)
-
-    def test_eviction_lru(self):
-        cache = Cache("test", max_entries=2)
-
-        cache.prefill(1, "one")
-        cache.prefill(2, "two")
-
-        # Now access 1 again, thus causing 2 to be least-recently used
-        cache.get(1)
-
-        cache.prefill(3, "three")
-
-        failed = False
-        try:
-            cache.get(2)
-        except KeyError:
-            failed = True
-
-        self.assertTrue(failed)
-
-        cache.get(1)
-        cache.get(3)
-
-
-class CacheDecoratorTestCase(unittest.HomeserverTestCase):
-    @defer.inlineCallbacks
-    def test_passthrough(self):
-        class A:
-            @cached()
-            def func(self, key):
-                return key
-
-        a = A()
-
-        self.assertEquals((yield a.func("foo")), "foo")
-        self.assertEquals((yield a.func("bar")), "bar")
-
-    @defer.inlineCallbacks
-    def test_hit(self):
-        callcount = [0]
-
-        class A:
-            @cached()
-            def func(self, key):
-                callcount[0] += 1
-                return key
-
-        a = A()
-        yield a.func("foo")
-
-        self.assertEquals(callcount[0], 1)
-
-        self.assertEquals((yield a.func("foo")), "foo")
-        self.assertEquals(callcount[0], 1)
-
-    @defer.inlineCallbacks
-    def test_invalidate(self):
-        callcount = [0]
-
-        class A:
-            @cached()
-            def func(self, key):
-                callcount[0] += 1
-                return key
-
-        a = A()
-        yield a.func("foo")
-
-        self.assertEquals(callcount[0], 1)
-
-        a.func.invalidate(("foo",))
-
-        yield a.func("foo")
-
-        self.assertEquals(callcount[0], 2)
-
-    def test_invalidate_missing(self):
-        class A:
-            @cached()
-            def func(self, key):
-                return key
-
-        A().func.invalidate(("what",))
-
-    @defer.inlineCallbacks
-    def test_max_entries(self):
-        callcount = [0]
-
-        class A:
-            @cached(max_entries=10)
-            def func(self, key):
-                callcount[0] += 1
-                return key
-
-        a = A()
-
-        for k in range(0, 12):
-            yield a.func(k)
-
-        self.assertEquals(callcount[0], 12)
-
-        # There must have been at least 2 evictions, meaning if we calculate
-        # all 12 values again, we must get called at least 2 more times
-        for k in range(0, 12):
-            yield a.func(k)
-
-        self.assertTrue(
-            callcount[0] >= 14, msg="Expected callcount >= 14, got %d" % (callcount[0])
-        )
-
-    def test_prefill(self):
-        callcount = [0]
-
-        d = defer.succeed(123)
-
-        class A:
-            @cached()
-            def func(self, key):
-                callcount[0] += 1
-                return d
-
-        a = A()
-
-        a.func.prefill(("foo",), ObservableDeferred(d))
-
-        self.assertEquals(a.func("foo").result, d.result)
-        self.assertEquals(callcount[0], 0)
-
-    @defer.inlineCallbacks
-    def test_invalidate_context(self):
-        callcount = [0]
-        callcount2 = [0]
-
-        class A:
-            @cached()
-            def func(self, key):
-                callcount[0] += 1
-                return key
-
-            @cached(cache_context=True)
-            def func2(self, key, cache_context):
-                callcount2[0] += 1
-                return self.func(key, on_invalidate=cache_context.invalidate)
-
-        a = A()
-        yield a.func2("foo")
-
-        self.assertEquals(callcount[0], 1)
-        self.assertEquals(callcount2[0], 1)
-
-        a.func.invalidate(("foo",))
-        yield a.func("foo")
-
-        self.assertEquals(callcount[0], 2)
-        self.assertEquals(callcount2[0], 1)
-
-        yield a.func2("foo")
-
-        self.assertEquals(callcount[0], 2)
-        self.assertEquals(callcount2[0], 2)
-
-    @defer.inlineCallbacks
-    def test_eviction_context(self):
-        callcount = [0]
-        callcount2 = [0]
-
-        class A:
-            @cached(max_entries=2)
-            def func(self, key):
-                callcount[0] += 1
-                return key
-
-            @cached(cache_context=True)
-            def func2(self, key, cache_context):
-                callcount2[0] += 1
-                return self.func(key, on_invalidate=cache_context.invalidate)
-
-        a = A()
-        yield a.func2("foo")
-        yield a.func2("foo2")
-
-        self.assertEquals(callcount[0], 2)
-        self.assertEquals(callcount2[0], 2)
-
-        yield a.func2("foo")
-        self.assertEquals(callcount[0], 2)
-        self.assertEquals(callcount2[0], 2)
-
-        yield a.func("foo3")
-
-        self.assertEquals(callcount[0], 3)
-        self.assertEquals(callcount2[0], 2)
-
-        yield a.func2("foo")
-
-        self.assertEquals(callcount[0], 4)
-        self.assertEquals(callcount2[0], 3)
-
-    @defer.inlineCallbacks
-    def test_double_get(self):
-        callcount = [0]
-        callcount2 = [0]
-
-        class A:
-            @cached()
-            def func(self, key):
-                callcount[0] += 1
-                return key
-
-            @cached(cache_context=True)
-            def func2(self, key, cache_context):
-                callcount2[0] += 1
-                return self.func(key, on_invalidate=cache_context.invalidate)
-
-        a = A()
-        a.func2.cache.cache = Mock(wraps=a.func2.cache.cache)
-
-        yield a.func2("foo")
-
-        self.assertEquals(callcount[0], 1)
-        self.assertEquals(callcount2[0], 1)
-
-        a.func2.invalidate(("foo",))
-        self.assertEquals(a.func2.cache.cache.pop.call_count, 1)
-
-        yield a.func2("foo")
-        a.func2.invalidate(("foo",))
-        self.assertEquals(a.func2.cache.cache.pop.call_count, 2)
-
-        self.assertEquals(callcount[0], 1)
-        self.assertEquals(callcount2[0], 2)
-
-        a.func.invalidate(("foo",))
-        self.assertEquals(a.func2.cache.cache.pop.call_count, 3)
-        yield a.func("foo")
-
-        self.assertEquals(callcount[0], 2)
-        self.assertEquals(callcount2[0], 2)
-
-        yield a.func2("foo")
-
-        self.assertEquals(callcount[0], 2)
-        self.assertEquals(callcount2[0], 3)
-
-
 class UpsertManyTests(unittest.HomeserverTestCase):
     def prepare(self, reactor, clock, hs):
         self.storage = hs.get_datastore()
diff --git a/tests/storage/test_appservice.py b/tests/storage/test_appservice.py
index 46f94914..1ce29af5 100644
--- a/tests/storage/test_appservice.py
+++ b/tests/storage/test_appservice.py
@@ -58,7 +58,7 @@ class ApplicationServiceStoreTestCase(unittest.TestCase):
         # must be done after inserts
         database = hs.get_datastores().databases[0]
         self.store = ApplicationServiceStore(
-            database, make_conn(database._database_config, database.engine), hs
+            database, make_conn(database._database_config, database.engine, "test"), hs
         )
 
     def tearDown(self):
@@ -132,7 +132,7 @@ class ApplicationServiceTransactionStoreTestCase(unittest.TestCase):
 
         db_config = hs.config.get_single_database()
         self.store = TestTransactionStore(
-            database, make_conn(db_config, self.engine), hs
+            database, make_conn(db_config, self.engine, "test"), hs
         )
 
     def _add_service(self, url, as_token, id):
@@ -244,7 +244,7 @@ class ApplicationServiceTransactionStoreTestCase(unittest.TestCase):
         service = Mock(id=self.as_list[0]["id"])
         events = [Mock(event_id="e1"), Mock(event_id="e2")]
         txn = yield defer.ensureDeferred(
-            self.store.create_appservice_txn(service, events)
+            self.store.create_appservice_txn(service, events, [])
         )
         self.assertEquals(txn.id, 1)
         self.assertEquals(txn.events, events)
@@ -258,7 +258,7 @@ class ApplicationServiceTransactionStoreTestCase(unittest.TestCase):
         yield self._insert_txn(service.id, 9644, events)
         yield self._insert_txn(service.id, 9645, events)
         txn = yield defer.ensureDeferred(
-            self.store.create_appservice_txn(service, events)
+            self.store.create_appservice_txn(service, events, [])
         )
         self.assertEquals(txn.id, 9646)
         self.assertEquals(txn.events, events)
@@ -270,7 +270,7 @@ class ApplicationServiceTransactionStoreTestCase(unittest.TestCase):
         events = [Mock(event_id="e1"), Mock(event_id="e2")]
         yield self._set_last_txn(service.id, 9643)
         txn = yield defer.ensureDeferred(
-            self.store.create_appservice_txn(service, events)
+            self.store.create_appservice_txn(service, events, [])
         )
         self.assertEquals(txn.id, 9644)
         self.assertEquals(txn.events, events)
@@ -293,7 +293,7 @@ class ApplicationServiceTransactionStoreTestCase(unittest.TestCase):
         yield self._insert_txn(self.as_list[3]["id"], 9643, events)
 
         txn = yield defer.ensureDeferred(
-            self.store.create_appservice_txn(service, events)
+            self.store.create_appservice_txn(service, events, [])
         )
         self.assertEquals(txn.id, 9644)
         self.assertEquals(txn.events, events)
@@ -410,6 +410,62 @@ class ApplicationServiceTransactionStoreTestCase(unittest.TestCase):
         )
 
 
+class ApplicationServiceStoreTypeStreamIds(unittest.HomeserverTestCase):
+    def make_homeserver(self, reactor, clock):
+        hs = self.setup_test_homeserver()
+        return hs
+
+    def prepare(self, hs, reactor, clock):
+        self.service = Mock(id="foo")
+        self.store = self.hs.get_datastore()
+        self.get_success(self.store.set_appservice_state(self.service, "up"))
+
+    def test_get_type_stream_id_for_appservice_no_value(self):
+        value = self.get_success(
+            self.store.get_type_stream_id_for_appservice(self.service, "read_receipt")
+        )
+        self.assertEquals(value, 0)
+
+        value = self.get_success(
+            self.store.get_type_stream_id_for_appservice(self.service, "presence")
+        )
+        self.assertEquals(value, 0)
+
+    def test_get_type_stream_id_for_appservice_invalid_type(self):
+        self.get_failure(
+            self.store.get_type_stream_id_for_appservice(self.service, "foobar"),
+            ValueError,
+        )
+
+    def test_set_type_stream_id_for_appservice(self):
+        read_receipt_value = 1024
+        self.get_success(
+            self.store.set_type_stream_id_for_appservice(
+                self.service, "read_receipt", read_receipt_value
+            )
+        )
+        result = self.get_success(
+            self.store.get_type_stream_id_for_appservice(self.service, "read_receipt")
+        )
+        self.assertEqual(result, read_receipt_value)
+
+        self.get_success(
+            self.store.set_type_stream_id_for_appservice(
+                self.service, "presence", read_receipt_value
+            )
+        )
+        result = self.get_success(
+            self.store.get_type_stream_id_for_appservice(self.service, "presence")
+        )
+        self.assertEqual(result, read_receipt_value)
+
+    def test_set_type_stream_id_for_appservice_invalid_type(self):
+        self.get_failure(
+            self.store.set_type_stream_id_for_appservice(self.service, "foobar", 1024),
+            ValueError,
+        )
+
+
 # required for ApplicationServiceTransactionStoreTestCase tests
 class TestTransactionStore(ApplicationServiceTransactionStore, ApplicationServiceStore):
     def __init__(self, database: DatabasePool, db_conn, hs):
@@ -448,7 +504,7 @@ class ApplicationServiceStoreConfigTestCase(unittest.TestCase):
 
         database = hs.get_datastores().databases[0]
         ApplicationServiceStore(
-            database, make_conn(database._database_config, database.engine), hs
+            database, make_conn(database._database_config, database.engine, "test"), hs
         )
 
     @defer.inlineCallbacks
@@ -467,7 +523,9 @@ class ApplicationServiceStoreConfigTestCase(unittest.TestCase):
         with self.assertRaises(ConfigError) as cm:
             database = hs.get_datastores().databases[0]
             ApplicationServiceStore(
-                database, make_conn(database._database_config, database.engine), hs
+                database,
+                make_conn(database._database_config, database.engine, "test"),
+                hs,
             )
 
         e = cm.exception
@@ -491,7 +549,9 @@ class ApplicationServiceStoreConfigTestCase(unittest.TestCase):
         with self.assertRaises(ConfigError) as cm:
             database = hs.get_datastores().databases[0]
             ApplicationServiceStore(
-                database, make_conn(database._database_config, database.engine), hs
+                database,
+                make_conn(database._database_config, database.engine, "test"),
+                hs,
             )
 
         e = cm.exception
diff --git a/tests/storage/test_id_generators.py b/tests/storage/test_id_generators.py
index 392b0883..cc0612cf 100644
--- a/tests/storage/test_id_generators.py
+++ b/tests/storage/test_id_generators.py
@@ -199,10 +199,17 @@ class MultiWriterIdGeneratorTestCase(HomeserverTestCase):
         first_id_gen = self._create_id_generator("first", writers=["first", "second"])
         second_id_gen = self._create_id_generator("second", writers=["first", "second"])
 
-        self.assertEqual(first_id_gen.get_positions(), {"first": 3, "second": 7})
-        self.assertEqual(first_id_gen.get_current_token_for_writer("first"), 3)
+        # The first ID gen will notice that it can advance its token to 7 as it
+        # has no in progress writes...
+        self.assertEqual(first_id_gen.get_positions(), {"first": 7, "second": 7})
+        self.assertEqual(first_id_gen.get_current_token_for_writer("first"), 7)
         self.assertEqual(first_id_gen.get_current_token_for_writer("second"), 7)
 
+        # ... but the second ID gen doesn't know that.
+        self.assertEqual(second_id_gen.get_positions(), {"first": 3, "second": 7})
+        self.assertEqual(second_id_gen.get_current_token_for_writer("first"), 3)
+        self.assertEqual(second_id_gen.get_current_token_for_writer("second"), 7)
+
         # Try allocating a new ID gen and check that we only see position
         # advanced after we leave the context manager.
 
@@ -211,7 +218,7 @@ class MultiWriterIdGeneratorTestCase(HomeserverTestCase):
                 self.assertEqual(stream_id, 8)
 
                 self.assertEqual(
-                    first_id_gen.get_positions(), {"first": 3, "second": 7}
+                    first_id_gen.get_positions(), {"first": 7, "second": 7}
                 )
 
         self.get_success(_get_next_async())
@@ -279,7 +286,7 @@ class MultiWriterIdGeneratorTestCase(HomeserverTestCase):
         self._insert_row_with_id("first", 3)
         self._insert_row_with_id("second", 5)
 
-        id_gen = self._create_id_generator("first", writers=["first", "second"])
+        id_gen = self._create_id_generator("worker", writers=["first", "second"])
 
         self.assertEqual(id_gen.get_positions(), {"first": 3, "second": 5})
 
@@ -319,14 +326,14 @@ class MultiWriterIdGeneratorTestCase(HomeserverTestCase):
 
         id_gen = self._create_id_generator("first", writers=["first", "second"])
 
-        self.assertEqual(id_gen.get_positions(), {"first": 3, "second": 5})
+        self.assertEqual(id_gen.get_positions(), {"first": 5, "second": 5})
 
-        self.assertEqual(id_gen.get_persisted_upto_position(), 3)
+        self.assertEqual(id_gen.get_persisted_upto_position(), 5)
 
         async def _get_next_async():
             async with id_gen.get_next() as stream_id:
                 self.assertEqual(stream_id, 6)
-                self.assertEqual(id_gen.get_persisted_upto_position(), 3)
+                self.assertEqual(id_gen.get_persisted_upto_position(), 5)
 
         self.get_success(_get_next_async())
 
@@ -388,7 +395,7 @@ class MultiWriterIdGeneratorTestCase(HomeserverTestCase):
         self._insert_row_with_id("second", 5)
 
         # Initial config has two writers
-        id_gen = self._create_id_generator("first", writers=["first", "second"])
+        id_gen = self._create_id_generator("worker", writers=["first", "second"])
         self.assertEqual(id_gen.get_persisted_upto_position(), 3)
         self.assertEqual(id_gen.get_current_token_for_writer("first"), 3)
         self.assertEqual(id_gen.get_current_token_for_writer("second"), 5)
@@ -568,7 +575,7 @@ class BackwardsMultiWriterIdGeneratorTestCase(HomeserverTestCase):
 
         self.get_success(_get_next_async2())
 
-        self.assertEqual(id_gen_1.get_positions(), {"first": -1, "second": -2})
+        self.assertEqual(id_gen_1.get_positions(), {"first": -2, "second": -2})
         self.assertEqual(id_gen_2.get_positions(), {"first": -1, "second": -2})
         self.assertEqual(id_gen_1.get_persisted_upto_position(), -2)
         self.assertEqual(id_gen_2.get_persisted_upto_position(), -2)
diff --git a/tests/storage/test_redaction.py b/tests/storage/test_redaction.py
index 1ea35d60..d4f9e809 100644
--- a/tests/storage/test_redaction.py
+++ b/tests/storage/test_redaction.py
@@ -236,9 +236,9 @@ class RedactionTestCase(unittest.HomeserverTestCase):
                 self._event_id = event_id
 
             @defer.inlineCallbacks
-            def build(self, prev_event_ids):
+            def build(self, prev_event_ids, auth_event_ids):
                 built_event = yield defer.ensureDeferred(
-                    self._base_builder.build(prev_event_ids)
+                    self._base_builder.build(prev_event_ids, auth_event_ids)
                 )
 
                 built_event._event_id = self._event_id
diff --git a/tests/test_federation.py b/tests/test_federation.py
index 27a7fc9e..d39e7925 100644
--- a/tests/test_federation.py
+++ b/tests/test_federation.py
@@ -75,7 +75,7 @@ class MessageAcceptTests(unittest.HomeserverTestCase):
             }
         )
 
-        self.handler = self.homeserver.get_handlers().federation_handler
+        self.handler = self.homeserver.get_federation_handler()
         self.handler.do_auth = lambda origin, event, context, auth_events: succeed(
             context
         )
diff --git a/tests/test_metrics.py b/tests/test_metrics.py
index f5f63d8e..759e4cd0 100644
--- a/tests/test_metrics.py
+++ b/tests/test_metrics.py
@@ -15,7 +15,7 @@
 # limitations under the License.
 
 from synapse.metrics import REGISTRY, InFlightGauge, generate_latest
-from synapse.util.caches.descriptors import Cache
+from synapse.util.caches.deferred_cache import DeferredCache
 
 from tests import unittest
 
@@ -138,7 +138,7 @@ class CacheMetricsTests(unittest.HomeserverTestCase):
         Caches produce metrics reflecting their state when scraped.
         """
         CACHE_NAME = "cache_metrics_test_fgjkbdfg"
-        cache = Cache(CACHE_NAME, max_entries=777)
+        cache = DeferredCache(CACHE_NAME, max_entries=777)
 
         items = {
             x.split(b"{")[0].decode("ascii"): x.split(b" ")[1].decode("ascii")
diff --git a/tests/test_phone_home.py b/tests/test_phone_home.py
index 7657bdde..e7aed092 100644
--- a/tests/test_phone_home.py
+++ b/tests/test_phone_home.py
@@ -17,7 +17,7 @@ import resource
 
 import mock
 
-from synapse.app.homeserver import phone_stats_home
+from synapse.app.phone_stats_home import phone_stats_home
 
 from tests.unittest import HomeserverTestCase
 
diff --git a/tests/unittest.py b/tests/unittest.py
index e654c044..040b126a 100644
--- a/tests/unittest.py
+++ b/tests/unittest.py
@@ -20,7 +20,7 @@ import hmac
 import inspect
 import logging
 import time
-from typing import Optional, Tuple, Type, TypeVar, Union
+from typing import Optional, Tuple, Type, TypeVar, Union, overload
 
 from mock import Mock, patch
 
@@ -241,7 +241,7 @@ class HomeserverTestCase(TestCase):
         # create a site to wrap the resource.
         self.site = SynapseSite(
             logger_name="synapse.access.http.fake",
-            site_tag="test",
+            site_tag=self.hs.config.server.server_name,
             config=self.hs.config.server.listeners[0],
             resource=self.resource,
             server_version_string="1",
@@ -254,17 +254,24 @@ class HomeserverTestCase(TestCase):
         if hasattr(self, "user_id"):
             if self.hijack_auth:
 
+                # We need a valid token ID to satisfy foreign key constraints.
+                token_id = self.get_success(
+                    self.hs.get_datastore().add_access_token_to_user(
+                        self.helper.auth_user_id, "some_fake_token", None, None,
+                    )
+                )
+
                 async def get_user_by_access_token(token=None, allow_guest=False):
                     return {
                         "user": UserID.from_string(self.helper.auth_user_id),
-                        "token_id": 1,
+                        "token_id": token_id,
                         "is_guest": False,
                     }
 
                 async def get_user_by_req(request, allow_guest=False, rights="access"):
                     return create_requester(
                         UserID.from_string(self.helper.auth_user_id),
-                        1,
+                        token_id,
                         False,
                         False,
                         None,
@@ -357,6 +364,36 @@ class HomeserverTestCase(TestCase):
         Function to optionally be overridden in subclasses.
         """
 
+    # Annoyingly mypy doesn't seem to pick up the fact that T is SynapseRequest
+    # when the `request` arg isn't given, so we define an explicit override to
+    # cover that case.
+    @overload
+    def make_request(
+        self,
+        method: Union[bytes, str],
+        path: Union[bytes, str],
+        content: Union[bytes, dict] = b"",
+        access_token: Optional[str] = None,
+        shorthand: bool = True,
+        federation_auth_origin: str = None,
+        content_is_form: bool = False,
+    ) -> Tuple[SynapseRequest, FakeChannel]:
+        ...
+
+    @overload
+    def make_request(
+        self,
+        method: Union[bytes, str],
+        path: Union[bytes, str],
+        content: Union[bytes, dict] = b"",
+        access_token: Optional[str] = None,
+        request: Type[T] = SynapseRequest,
+        shorthand: bool = True,
+        federation_auth_origin: str = None,
+        content_is_form: bool = False,
+    ) -> Tuple[T, FakeChannel]:
+        ...
+
     def make_request(
         self,
         method: Union[bytes, str],
@@ -608,7 +645,9 @@ class HomeserverTestCase(TestCase):
         if soft_failed:
             event.internal_metadata.soft_failed = True
 
-        self.get_success(event_creator.send_nonmember_event(requester, event, context))
+        self.get_success(
+            event_creator.handle_new_client_event(requester, event, context)
+        )
 
         return event.event_id
 
diff --git a/tests/util/caches/test_deferred_cache.py b/tests/util/caches/test_deferred_cache.py
new file mode 100644
index 00000000..dadfabd4
--- /dev/null
+++ b/tests/util/caches/test_deferred_cache.py
@@ -0,0 +1,251 @@
+# -*- coding: utf-8 -*-
+# Copyright 2020 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 functools import partial
+
+from twisted.internet import defer
+
+from synapse.util.caches.deferred_cache import DeferredCache
+
+from tests.unittest import TestCase
+
+
+class DeferredCacheTestCase(TestCase):
+    def test_empty(self):
+        cache = DeferredCache("test")
+        failed = False
+        try:
+            cache.get("foo")
+        except KeyError:
+            failed = True
+
+        self.assertTrue(failed)
+
+    def test_hit(self):
+        cache = DeferredCache("test")
+        cache.prefill("foo", 123)
+
+        self.assertEquals(self.successResultOf(cache.get("foo")), 123)
+
+    def test_hit_deferred(self):
+        cache = DeferredCache("test")
+        origin_d = defer.Deferred()
+        set_d = cache.set("k1", origin_d)
+
+        # get should return an incomplete deferred
+        get_d = cache.get("k1")
+        self.assertFalse(get_d.called)
+
+        # add a callback that will make sure that the set_d gets called before the get_d
+        def check1(r):
+            self.assertTrue(set_d.called)
+            return r
+
+        # TODO: Actually ObservableDeferred *doesn't* run its tests in order on py3.8.
+        #   maybe we should fix that?
+        # get_d.addCallback(check1)
+
+        # now fire off all the deferreds
+        origin_d.callback(99)
+        self.assertEqual(self.successResultOf(origin_d), 99)
+        self.assertEqual(self.successResultOf(set_d), 99)
+        self.assertEqual(self.successResultOf(get_d), 99)
+
+    def test_callbacks(self):
+        """Invalidation callbacks are called at the right time"""
+        cache = DeferredCache("test")
+        callbacks = set()
+
+        # start with an entry, with a callback
+        cache.prefill("k1", 10, callback=lambda: callbacks.add("prefill"))
+
+        # now replace that entry with a pending result
+        origin_d = defer.Deferred()
+        set_d = cache.set("k1", origin_d, callback=lambda: callbacks.add("set"))
+
+        # ... and also make a get request
+        get_d = cache.get("k1", callback=lambda: callbacks.add("get"))
+
+        # we don't expect the invalidation callback for the original value to have
+        # been called yet, even though get() will now return a different result.
+        # I'm not sure if that is by design or not.
+        self.assertEqual(callbacks, set())
+
+        # now fire off all the deferreds
+        origin_d.callback(20)
+        self.assertEqual(self.successResultOf(set_d), 20)
+        self.assertEqual(self.successResultOf(get_d), 20)
+
+        # now the original invalidation callback should have been called, but none of
+        # the others
+        self.assertEqual(callbacks, {"prefill"})
+        callbacks.clear()
+
+        # another update should invalidate both the previous results
+        cache.prefill("k1", 30)
+        self.assertEqual(callbacks, {"set", "get"})
+
+    def test_set_fail(self):
+        cache = DeferredCache("test")
+        callbacks = set()
+
+        # start with an entry, with a callback
+        cache.prefill("k1", 10, callback=lambda: callbacks.add("prefill"))
+
+        # now replace that entry with a pending result
+        origin_d = defer.Deferred()
+        set_d = cache.set("k1", origin_d, callback=lambda: callbacks.add("set"))
+
+        # ... and also make a get request
+        get_d = cache.get("k1", callback=lambda: callbacks.add("get"))
+
+        # none of the callbacks should have been called yet
+        self.assertEqual(callbacks, set())
+
+        # oh noes! fails!
+        e = Exception("oops")
+        origin_d.errback(e)
+        self.assertIs(self.failureResultOf(set_d, Exception).value, e)
+        self.assertIs(self.failureResultOf(get_d, Exception).value, e)
+
+        # the callbacks for the failed requests should have been called.
+        # I'm not sure if this is deliberate or not.
+        self.assertEqual(callbacks, {"get", "set"})
+        callbacks.clear()
+
+        # the old value should still be returned now?
+        get_d2 = cache.get("k1", callback=lambda: callbacks.add("get2"))
+        self.assertEqual(self.successResultOf(get_d2), 10)
+
+        # replacing the value now should run the callbacks for those requests
+        # which got the original result
+        cache.prefill("k1", 30)
+        self.assertEqual(callbacks, {"prefill", "get2"})
+
+    def test_get_immediate(self):
+        cache = DeferredCache("test")
+        d1 = defer.Deferred()
+        cache.set("key1", d1)
+
+        # get_immediate should return default
+        v = cache.get_immediate("key1", 1)
+        self.assertEqual(v, 1)
+
+        # now complete the set
+        d1.callback(2)
+
+        # get_immediate should return result
+        v = cache.get_immediate("key1", 1)
+        self.assertEqual(v, 2)
+
+    def test_invalidate(self):
+        cache = DeferredCache("test")
+        cache.prefill(("foo",), 123)
+        cache.invalidate(("foo",))
+
+        failed = False
+        try:
+            cache.get(("foo",))
+        except KeyError:
+            failed = True
+
+        self.assertTrue(failed)
+
+    def test_invalidate_all(self):
+        cache = DeferredCache("testcache")
+
+        callback_record = [False, False]
+
+        def record_callback(idx):
+            callback_record[idx] = True
+
+        # add a couple of pending entries
+        d1 = defer.Deferred()
+        cache.set("key1", d1, partial(record_callback, 0))
+
+        d2 = defer.Deferred()
+        cache.set("key2", d2, partial(record_callback, 1))
+
+        # lookup should return pending deferreds
+        self.assertFalse(cache.get("key1").called)
+        self.assertFalse(cache.get("key2").called)
+
+        # let one of the lookups complete
+        d2.callback("result2")
+
+        # now the cache will return a completed deferred
+        self.assertEqual(self.successResultOf(cache.get("key2")), "result2")
+
+        # now do the invalidation
+        cache.invalidate_all()
+
+        # lookup should fail
+        with self.assertRaises(KeyError):
+            cache.get("key1")
+        with self.assertRaises(KeyError):
+            cache.get("key2")
+
+        # both callbacks should have been callbacked
+        self.assertTrue(callback_record[0], "Invalidation callback for key1 not called")
+        self.assertTrue(callback_record[1], "Invalidation callback for key2 not called")
+
+        # letting the other lookup complete should do nothing
+        d1.callback("result1")
+        with self.assertRaises(KeyError):
+            cache.get("key1", None)
+
+    def test_eviction(self):
+        cache = DeferredCache(
+            "test", max_entries=2, apply_cache_factor_from_config=False
+        )
+
+        cache.prefill(1, "one")
+        cache.prefill(2, "two")
+        cache.prefill(3, "three")  # 1 will be evicted
+
+        failed = False
+        try:
+            cache.get(1)
+        except KeyError:
+            failed = True
+
+        self.assertTrue(failed)
+
+        cache.get(2)
+        cache.get(3)
+
+    def test_eviction_lru(self):
+        cache = DeferredCache(
+            "test", max_entries=2, apply_cache_factor_from_config=False
+        )
+
+        cache.prefill(1, "one")
+        cache.prefill(2, "two")
+
+        # Now access 1 again, thus causing 2 to be least-recently used
+        cache.get(1)
+
+        cache.prefill(3, "three")
+
+        failed = False
+        try:
+            cache.get(2)
+        except KeyError:
+            failed = True
+
+        self.assertTrue(failed)
+
+        cache.get(1)
+        cache.get(3)
diff --git a/tests/util/caches/test_descriptors.py b/tests/util/caches/test_descriptors.py
index 677e9254..2ad08f54 100644
--- a/tests/util/caches/test_descriptors.py
+++ b/tests/util/caches/test_descriptors.py
@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
-from functools import partial
+from typing import Set
 
 import mock
 
@@ -42,49 +42,6 @@ def run_on_reactor():
     return make_deferred_yieldable(d)
 
 
-class CacheTestCase(unittest.TestCase):
-    def test_invalidate_all(self):
-        cache = descriptors.Cache("testcache")
-
-        callback_record = [False, False]
-
-        def record_callback(idx):
-            callback_record[idx] = True
-
-        # add a couple of pending entries
-        d1 = defer.Deferred()
-        cache.set("key1", d1, partial(record_callback, 0))
-
-        d2 = defer.Deferred()
-        cache.set("key2", d2, partial(record_callback, 1))
-
-        # lookup should return observable deferreds
-        self.assertFalse(cache.get("key1").has_called())
-        self.assertFalse(cache.get("key2").has_called())
-
-        # let one of the lookups complete
-        d2.callback("result2")
-
-        # for now at least, the cache will return real results rather than an
-        # observabledeferred
-        self.assertEqual(cache.get("key2"), "result2")
-
-        # now do the invalidation
-        cache.invalidate_all()
-
-        # lookup should return none
-        self.assertIsNone(cache.get("key1", None))
-        self.assertIsNone(cache.get("key2", None))
-
-        # both callbacks should have been callbacked
-        self.assertTrue(callback_record[0], "Invalidation callback for key1 not called")
-        self.assertTrue(callback_record[1], "Invalidation callback for key2 not called")
-
-        # letting the other lookup complete should do nothing
-        d1.callback("result1")
-        self.assertIsNone(cache.get("key1", None))
-
-
 class DescriptorTestCase(unittest.TestCase):
     @defer.inlineCallbacks
     def test_cache(self):
@@ -174,6 +131,57 @@ class DescriptorTestCase(unittest.TestCase):
         d = obj.fn(1)
         self.failureResultOf(d, SynapseError)
 
+    def test_cache_with_async_exception(self):
+        """The wrapped function returns a failure
+        """
+
+        class Cls:
+            result = None
+            call_count = 0
+
+            @cached()
+            def fn(self, arg1):
+                self.call_count += 1
+                return self.result
+
+        obj = Cls()
+        callbacks = set()  # type: Set[str]
+
+        # set off an asynchronous request
+        obj.result = origin_d = defer.Deferred()
+
+        d1 = obj.fn(1, on_invalidate=lambda: callbacks.add("d1"))
+        self.assertFalse(d1.called)
+
+        # a second request should also return a deferred, but should not call the
+        # function itself.
+        d2 = obj.fn(1, on_invalidate=lambda: callbacks.add("d2"))
+        self.assertFalse(d2.called)
+        self.assertEqual(obj.call_count, 1)
+
+        # no callbacks yet
+        self.assertEqual(callbacks, set())
+
+        # the original request fails
+        e = Exception("bzz")
+        origin_d.errback(e)
+
+        # ... which should cause the lookups to fail similarly
+        self.assertIs(self.failureResultOf(d1, Exception).value, e)
+        self.assertIs(self.failureResultOf(d2, Exception).value, e)
+
+        # ... and the callbacks to have been, uh, called.
+        self.assertEqual(callbacks, {"d1", "d2"})
+
+        # ... leaving the cache empty
+        self.assertEqual(len(obj.fn.cache.cache), 0)
+
+        # and a second call should work as normal
+        obj.result = defer.succeed(100)
+        d3 = obj.fn(1)
+        self.assertEqual(self.successResultOf(d3), 100)
+        self.assertEqual(obj.call_count, 2)
+
     def test_cache_logcontexts(self):
         """Check that logcontexts are set and restored correctly when
         using the cache."""
@@ -355,6 +363,235 @@ class DescriptorTestCase(unittest.TestCase):
         self.failureResultOf(d, SynapseError)
 
 
+class CacheDecoratorTestCase(unittest.HomeserverTestCase):
+    """More tests for @cached
+
+    The following is a set of tests that got lost in a different file for a while.
+
+    There are probably duplicates of the tests in DescriptorTestCase. Ideally the
+    duplicates would be removed and the two sets of classes combined.
+    """
+
+    @defer.inlineCallbacks
+    def test_passthrough(self):
+        class A:
+            @cached()
+            def func(self, key):
+                return key
+
+        a = A()
+
+        self.assertEquals((yield a.func("foo")), "foo")
+        self.assertEquals((yield a.func("bar")), "bar")
+
+    @defer.inlineCallbacks
+    def test_hit(self):
+        callcount = [0]
+
+        class A:
+            @cached()
+            def func(self, key):
+                callcount[0] += 1
+                return key
+
+        a = A()
+        yield a.func("foo")
+
+        self.assertEquals(callcount[0], 1)
+
+        self.assertEquals((yield a.func("foo")), "foo")
+        self.assertEquals(callcount[0], 1)
+
+    @defer.inlineCallbacks
+    def test_invalidate(self):
+        callcount = [0]
+
+        class A:
+            @cached()
+            def func(self, key):
+                callcount[0] += 1
+                return key
+
+        a = A()
+        yield a.func("foo")
+
+        self.assertEquals(callcount[0], 1)
+
+        a.func.invalidate(("foo",))
+
+        yield a.func("foo")
+
+        self.assertEquals(callcount[0], 2)
+
+    def test_invalidate_missing(self):
+        class A:
+            @cached()
+            def func(self, key):
+                return key
+
+        A().func.invalidate(("what",))
+
+    @defer.inlineCallbacks
+    def test_max_entries(self):
+        callcount = [0]
+
+        class A:
+            @cached(max_entries=10)
+            def func(self, key):
+                callcount[0] += 1
+                return key
+
+        a = A()
+
+        for k in range(0, 12):
+            yield a.func(k)
+
+        self.assertEquals(callcount[0], 12)
+
+        # There must have been at least 2 evictions, meaning if we calculate
+        # all 12 values again, we must get called at least 2 more times
+        for k in range(0, 12):
+            yield a.func(k)
+
+        self.assertTrue(
+            callcount[0] >= 14, msg="Expected callcount >= 14, got %d" % (callcount[0])
+        )
+
+    def test_prefill(self):
+        callcount = [0]
+
+        d = defer.succeed(123)
+
+        class A:
+            @cached()
+            def func(self, key):
+                callcount[0] += 1
+                return d
+
+        a = A()
+
+        a.func.prefill(("foo",), 456)
+
+        self.assertEquals(a.func("foo").result, 456)
+        self.assertEquals(callcount[0], 0)
+
+    @defer.inlineCallbacks
+    def test_invalidate_context(self):
+        callcount = [0]
+        callcount2 = [0]
+
+        class A:
+            @cached()
+            def func(self, key):
+                callcount[0] += 1
+                return key
+
+            @cached(cache_context=True)
+            def func2(self, key, cache_context):
+                callcount2[0] += 1
+                return self.func(key, on_invalidate=cache_context.invalidate)
+
+        a = A()
+        yield a.func2("foo")
+
+        self.assertEquals(callcount[0], 1)
+        self.assertEquals(callcount2[0], 1)
+
+        a.func.invalidate(("foo",))
+        yield a.func("foo")
+
+        self.assertEquals(callcount[0], 2)
+        self.assertEquals(callcount2[0], 1)
+
+        yield a.func2("foo")
+
+        self.assertEquals(callcount[0], 2)
+        self.assertEquals(callcount2[0], 2)
+
+    @defer.inlineCallbacks
+    def test_eviction_context(self):
+        callcount = [0]
+        callcount2 = [0]
+
+        class A:
+            @cached(max_entries=2)
+            def func(self, key):
+                callcount[0] += 1
+                return key
+
+            @cached(cache_context=True)
+            def func2(self, key, cache_context):
+                callcount2[0] += 1
+                return self.func(key, on_invalidate=cache_context.invalidate)
+
+        a = A()
+        yield a.func2("foo")
+        yield a.func2("foo2")
+
+        self.assertEquals(callcount[0], 2)
+        self.assertEquals(callcount2[0], 2)
+
+        yield a.func2("foo")
+        self.assertEquals(callcount[0], 2)
+        self.assertEquals(callcount2[0], 2)
+
+        yield a.func("foo3")
+
+        self.assertEquals(callcount[0], 3)
+        self.assertEquals(callcount2[0], 2)
+
+        yield a.func2("foo")
+
+        self.assertEquals(callcount[0], 4)
+        self.assertEquals(callcount2[0], 3)
+
+    @defer.inlineCallbacks
+    def test_double_get(self):
+        callcount = [0]
+        callcount2 = [0]
+
+        class A:
+            @cached()
+            def func(self, key):
+                callcount[0] += 1
+                return key
+
+            @cached(cache_context=True)
+            def func2(self, key, cache_context):
+                callcount2[0] += 1
+                return self.func(key, on_invalidate=cache_context.invalidate)
+
+        a = A()
+        a.func2.cache.cache = mock.Mock(wraps=a.func2.cache.cache)
+
+        yield a.func2("foo")
+
+        self.assertEquals(callcount[0], 1)
+        self.assertEquals(callcount2[0], 1)
+
+        a.func2.invalidate(("foo",))
+        self.assertEquals(a.func2.cache.cache.pop.call_count, 1)
+
+        yield a.func2("foo")
+        a.func2.invalidate(("foo",))
+        self.assertEquals(a.func2.cache.cache.pop.call_count, 2)
+
+        self.assertEquals(callcount[0], 1)
+        self.assertEquals(callcount2[0], 2)
+
+        a.func.invalidate(("foo",))
+        self.assertEquals(a.func2.cache.cache.pop.call_count, 3)
+        yield a.func("foo")
+
+        self.assertEquals(callcount[0], 2)
+        self.assertEquals(callcount2[0], 2)
+
+        yield a.func2("foo")
+
+        self.assertEquals(callcount[0], 2)
+        self.assertEquals(callcount2[0], 3)
+
+
 class CachedListDescriptorTestCase(unittest.TestCase):
     @defer.inlineCallbacks
     def test_cache(self):
diff --git a/tests/util/test_lrucache.py b/tests/util/test_lrucache.py
index 0adb2174..a739a6aa 100644
--- a/tests/util/test_lrucache.py
+++ b/tests/util/test_lrucache.py
@@ -19,7 +19,8 @@ from mock import Mock
 from synapse.util.caches.lrucache import LruCache
 from synapse.util.caches.treecache import TreeCache
 
-from .. import unittest
+from tests import unittest
+from tests.unittest import override_config
 
 
 class LruCacheTestCase(unittest.HomeserverTestCase):
@@ -59,7 +60,7 @@ class LruCacheTestCase(unittest.HomeserverTestCase):
         self.assertEquals(cache.pop("key"), None)
 
     def test_del_multi(self):
-        cache = LruCache(4, 2, cache_type=TreeCache)
+        cache = LruCache(4, keylen=2, cache_type=TreeCache)
         cache[("animal", "cat")] = "mew"
         cache[("animal", "dog")] = "woof"
         cache[("vehicles", "car")] = "vroom"
@@ -83,6 +84,11 @@ class LruCacheTestCase(unittest.HomeserverTestCase):
         cache.clear()
         self.assertEquals(len(cache), 0)
 
+    @override_config({"caches": {"per_cache_factors": {"mycache": 10}}})
+    def test_special_size(self):
+        cache = LruCache(10, "mycache")
+        self.assertEqual(cache.max_size, 100)
+
 
 class LruCacheCallbacksTestCase(unittest.HomeserverTestCase):
     def test_get(self):
@@ -160,7 +166,7 @@ class LruCacheCallbacksTestCase(unittest.HomeserverTestCase):
         m2 = Mock()
         m3 = Mock()
         m4 = Mock()
-        cache = LruCache(4, 2, cache_type=TreeCache)
+        cache = LruCache(4, keylen=2, cache_type=TreeCache)
 
         cache.set(("a", "1"), "value", callbacks=[m1])
         cache.set(("a", "2"), "value", callbacks=[m2])
diff --git a/tests/utils.py b/tests/utils.py
index 4673872f..acec74e9 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -21,6 +21,7 @@ import time
 import uuid
 import warnings
 from inspect import getcallargs
+from typing import Type
 from urllib import parse as urlparse
 
 from mock import Mock, patch
@@ -38,6 +39,7 @@ from synapse.http.server import HttpServer
 from synapse.logging.context import current_context, set_current_context
 from synapse.server import HomeServer
 from synapse.storage import DataStore
+from synapse.storage.database import LoggingDatabaseConnection
 from synapse.storage.engines import PostgresEngine, create_engine
 from synapse.storage.prepare_database import prepare_database
 from synapse.util.ratelimitutils import FederationRateLimiter
@@ -88,6 +90,7 @@ def setupdb():
             host=POSTGRES_HOST,
             password=POSTGRES_PASSWORD,
         )
+        db_conn = LoggingDatabaseConnection(db_conn, db_engine, "tests")
         prepare_database(db_conn, db_engine, None)
         db_conn.close()
 
@@ -190,11 +193,10 @@ class TestHomeServer(HomeServer):
 def setup_test_homeserver(
     cleanup_func,
     name="test",
-    datastore=None,
     config=None,
     reactor=None,
-    homeserverToUse=TestHomeServer,
-    **kargs
+    homeserver_to_use: Type[HomeServer] = TestHomeServer,
+    **kwargs
 ):
     """
     Setup a homeserver suitable for running tests against.  Keyword arguments
@@ -217,8 +219,8 @@ def setup_test_homeserver(
 
     config.ldap_enabled = False
 
-    if "clock" not in kargs:
-        kargs["clock"] = MockClock()
+    if "clock" not in kwargs:
+        kwargs["clock"] = MockClock()
 
     if USE_POSTGRES_FOR_TESTS:
         test_db = "synapse_test_%s" % uuid.uuid4().hex
@@ -247,7 +249,7 @@ def setup_test_homeserver(
 
     # Create the database before we actually try and connect to it, based off
     # the template database we generate in setupdb()
-    if datastore is None and isinstance(db_engine, PostgresEngine):
+    if isinstance(db_engine, PostgresEngine):
         db_conn = db_engine.module.connect(
             database=POSTGRES_BASE_DB,
             user=POSTGRES_USER,
@@ -263,79 +265,68 @@ def setup_test_homeserver(
         cur.close()
         db_conn.close()
 
-    if datastore is None:
-        hs = homeserverToUse(
-            name,
-            config=config,
-            version_string="Synapse/tests",
-            tls_server_context_factory=Mock(),
-            tls_client_options_factory=Mock(),
-            reactor=reactor,
-            **kargs
-        )
+    hs = homeserver_to_use(
+        name, config=config, version_string="Synapse/tests", reactor=reactor,
+    )
 
-        hs.setup()
-        if homeserverToUse.__name__ == "TestHomeServer":
-            hs.setup_master()
+    # Install @cache_in_self attributes
+    for key, val in kwargs.items():
+        setattr(hs, key, val)
 
-        if isinstance(db_engine, PostgresEngine):
-            database = hs.get_datastores().databases[0]
+    # Mock TLS
+    hs.tls_server_context_factory = Mock()
+    hs.tls_client_options_factory = Mock()
 
-            # We need to do cleanup on PostgreSQL
-            def cleanup():
-                import psycopg2
+    hs.setup()
+    if homeserver_to_use == TestHomeServer:
+        hs.setup_background_tasks()
 
-                # Close all the db pools
-                database._db_pool.close()
+    if isinstance(db_engine, PostgresEngine):
+        database = hs.get_datastores().databases[0]
 
-                dropped = False
+        # We need to do cleanup on PostgreSQL
+        def cleanup():
+            import psycopg2
 
-                # Drop the test database
-                db_conn = db_engine.module.connect(
-                    database=POSTGRES_BASE_DB,
-                    user=POSTGRES_USER,
-                    host=POSTGRES_HOST,
-                    password=POSTGRES_PASSWORD,
-                )
-                db_conn.autocommit = True
-                cur = db_conn.cursor()
-
-                # Try a few times to drop the DB. Some things may hold on to the
-                # database for a few more seconds due to flakiness, preventing
-                # us from dropping it when the test is over. If we can't drop
-                # it, warn and move on.
-                for x in range(5):
-                    try:
-                        cur.execute("DROP DATABASE IF EXISTS %s;" % (test_db,))
-                        db_conn.commit()
-                        dropped = True
-                    except psycopg2.OperationalError as e:
-                        warnings.warn(
-                            "Couldn't drop old db: " + str(e), category=UserWarning
-                        )
-                        time.sleep(0.5)
-
-                cur.close()
-                db_conn.close()
-
-                if not dropped:
-                    warnings.warn("Failed to drop old DB.", category=UserWarning)
-
-            if not LEAVE_DB:
-                # Register the cleanup hook
-                cleanup_func(cleanup)
+            # Close all the db pools
+            database._db_pool.close()
 
-    else:
-        hs = homeserverToUse(
-            name,
-            datastore=datastore,
-            config=config,
-            version_string="Synapse/tests",
-            tls_server_context_factory=Mock(),
-            tls_client_options_factory=Mock(),
-            reactor=reactor,
-            **kargs
-        )
+            dropped = False
+
+            # Drop the test database
+            db_conn = db_engine.module.connect(
+                database=POSTGRES_BASE_DB,
+                user=POSTGRES_USER,
+                host=POSTGRES_HOST,
+                password=POSTGRES_PASSWORD,
+            )
+            db_conn.autocommit = True
+            cur = db_conn.cursor()
+
+            # Try a few times to drop the DB. Some things may hold on to the
+            # database for a few more seconds due to flakiness, preventing
+            # us from dropping it when the test is over. If we can't drop
+            # it, warn and move on.
+            for x in range(5):
+                try:
+                    cur.execute("DROP DATABASE IF EXISTS %s;" % (test_db,))
+                    db_conn.commit()
+                    dropped = True
+                except psycopg2.OperationalError as e:
+                    warnings.warn(
+                        "Couldn't drop old db: " + str(e), category=UserWarning
+                    )
+                    time.sleep(0.5)
+
+            cur.close()
+            db_conn.close()
+
+            if not dropped:
+                warnings.warn("Failed to drop old DB.", category=UserWarning)
+
+        if not LEAVE_DB:
+            # Register the cleanup hook
+            cleanup_func(cleanup)
 
     # bcrypt is far too slow to be doing in unit tests
     # Need to let the HS build an auth handler and then mess with it
@@ -351,7 +342,7 @@ def setup_test_homeserver(
 
     hs.get_auth_handler().validate_hash = validate_hash
 
-    fed = kargs.get("resource_for_federation", None)
+    fed = kwargs.get("resource_for_federation", None)
     if fed:
         register_federation_servlets(hs, fed)