Move additional tasks to the background worker, part 4 (#8513)

1 files changed, 14 insertions, 34 deletions

diff --git a/tests/handlers/test_typing.py b/tests/handlers/test_typing.py
index 3fec09ea8a..16ff2e22d2 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(