Add type hints to tests files. (#12256)

1 files changed, 26 insertions, 22 deletions

diff --git a/tests/storage/test_background_update.py b/tests/storage/test_background_update.py
index 5cf18b690e..fd619b64d4 100644
--- a/tests/storage/test_background_update.py
+++ b/tests/storage/test_background_update.py
@@ -17,8 +17,12 @@ from unittest.mock import Mock
 import yaml
 
 from twisted.internet.defer import Deferred, ensureDeferred
+from twisted.test.proto_helpers import MemoryReactor
 
+from synapse.server import HomeServer
 from synapse.storage.background_updates import BackgroundUpdater
+from synapse.types import JsonDict
+from synapse.util import Clock
 
 from tests import unittest
 from tests.test_utils import make_awaitable, simple_async_mock
@@ -26,7 +30,7 @@ from tests.unittest import override_config
 
 
 class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
-    def prepare(self, reactor, clock, homeserver):
+    def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
         self.updates: BackgroundUpdater = self.hs.get_datastores().main.db_pool.updates
         # the base test class should have run the real bg updates for us
         self.assertTrue(
@@ -39,7 +43,7 @@ class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
         )
         self.store = self.hs.get_datastores().main
 
-    async def update(self, progress, count):
+    async def update(self, progress: JsonDict, count: int) -> int:
         duration_ms = 10
         await self.clock.sleep((count * duration_ms) / 1000)
         progress = {"my_key": progress["my_key"] + 1}
@@ -51,7 +55,7 @@ class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
         )
         return count
 
-    def test_do_background_update(self):
+    def test_do_background_update(self) -> None:
         # the time we claim it takes to update one item when running the update
         duration_ms = 10
 
@@ -80,7 +84,7 @@ class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
 
         # second step: complete the update
         # we should now get run with a much bigger number of items to update
-        async def update(progress, count):
+        async def update(progress: JsonDict, count: int) -> int:
             self.assertEqual(progress, {"my_key": 2})
             self.assertAlmostEqual(
                 count,
@@ -110,7 +114,7 @@ class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
             """
         )
     )
-    def test_background_update_default_batch_set_by_config(self):
+    def test_background_update_default_batch_set_by_config(self) -> None:
         """
         Test that the background update is run with the default_batch_size set by the config
         """
@@ -133,7 +137,7 @@ class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
         # on the first call, we should get run with the default background update size specified in the config
         self.update_handler.assert_called_once_with({"my_key": 1}, 20)
 
-    def test_background_update_default_sleep_behavior(self):
+    def test_background_update_default_sleep_behavior(self) -> None:
         """
         Test default background update behavior, which is to sleep
         """
@@ -147,7 +151,7 @@ class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
 
         self.update_handler.side_effect = self.update
         self.update_handler.reset_mock()
-        self.updates.start_doing_background_updates(),
+        self.updates.start_doing_background_updates()
 
         # 2: advance the reactor less than the default sleep duration (1000ms)
         self.reactor.pump([0.5])
@@ -167,7 +171,7 @@ class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
             """
         )
     )
-    def test_background_update_sleep_set_in_config(self):
+    def test_background_update_sleep_set_in_config(self) -> None:
         """
         Test that changing the sleep time in the config changes how long it sleeps
         """
@@ -181,7 +185,7 @@ class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
 
         self.update_handler.side_effect = self.update
         self.update_handler.reset_mock()
-        self.updates.start_doing_background_updates(),
+        self.updates.start_doing_background_updates()
 
         # 2: advance the reactor less than the configured sleep duration (500ms)
         self.reactor.pump([0.45])
@@ -201,7 +205,7 @@ class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
             """
         )
     )
-    def test_disabling_background_update_sleep(self):
+    def test_disabling_background_update_sleep(self) -> None:
         """
         Test that disabling sleep in the config results in bg update not sleeping
         """
@@ -215,7 +219,7 @@ class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
 
         self.update_handler.side_effect = self.update
         self.update_handler.reset_mock()
-        self.updates.start_doing_background_updates(),
+        self.updates.start_doing_background_updates()
 
         # 2: advance the reactor very little
         self.reactor.pump([0.025])
@@ -230,7 +234,7 @@ class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
             """
         )
     )
-    def test_background_update_duration_set_in_config(self):
+    def test_background_update_duration_set_in_config(self) -> None:
         """
         Test that the desired duration set in the config is used in determining batch size
         """
@@ -254,7 +258,7 @@ class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
 
         # the first update was run with the default batch size, this should be run with 500ms as the
         # desired duration
-        async def update(progress, count):
+        async def update(progress: JsonDict, count: int) -> int:
             self.assertEqual(progress, {"my_key": 2})
             self.assertAlmostEqual(
                 count,
@@ -275,7 +279,7 @@ class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
             """
         )
     )
-    def test_background_update_min_batch_set_in_config(self):
+    def test_background_update_min_batch_set_in_config(self) -> None:
         """
         Test that the minimum batch size set in the config is used
         """
@@ -290,7 +294,7 @@ class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
         )
 
         # Run the update with the long-running update item
-        async def update(progress, count):
+        async def update_long(progress: JsonDict, count: int) -> int:
             await self.clock.sleep((count * duration_ms) / 1000)
             progress = {"my_key": progress["my_key"] + 1}
             await self.store.db_pool.runInteraction(
@@ -301,7 +305,7 @@ class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
             )
             return count
 
-        self.update_handler.side_effect = update
+        self.update_handler.side_effect = update_long
         self.update_handler.reset_mock()
         res = self.get_success(
             self.updates.do_next_background_update(False),
@@ -311,25 +315,25 @@ class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
 
         # the first update was run with the default batch size, this should be run with minimum batch size
         # as the first items took a very long time
-        async def update(progress, count):
+        async def update_short(progress: JsonDict, count: int) -> int:
             self.assertEqual(progress, {"my_key": 2})
             self.assertEqual(count, 5)
             await self.updates._end_background_update("test_update")
             return count
 
-        self.update_handler.side_effect = update
+        self.update_handler.side_effect = update_short
         self.get_success(self.updates.do_next_background_update(False))
 
 
 class BackgroundUpdateControllerTestCase(unittest.HomeserverTestCase):
-    def prepare(self, reactor, clock, homeserver):
+    def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
         self.updates: BackgroundUpdater = self.hs.get_datastores().main.db_pool.updates
         # the base test class should have run the real bg updates for us
         self.assertTrue(
             self.get_success(self.updates.has_completed_background_updates())
         )
 
-        self.update_deferred = Deferred()
+        self.update_deferred: Deferred[int] = Deferred()
         self.update_handler = Mock(return_value=self.update_deferred)
         self.updates.register_background_update_handler(
             "test_update", self.update_handler
@@ -358,7 +362,7 @@ class BackgroundUpdateControllerTestCase(unittest.HomeserverTestCase):
             ),
         )
 
-    def test_controller(self):
+    def test_controller(self) -> None:
         store = self.hs.get_datastores().main
         self.get_success(
             store.db_pool.simple_insert(
@@ -368,7 +372,7 @@ class BackgroundUpdateControllerTestCase(unittest.HomeserverTestCase):
         )
 
         # Set the return value for the context manager.
-        enter_defer = Deferred()
+        enter_defer: Deferred[int] = Deferred()
         self._update_ctx_manager.__aenter__ = Mock(return_value=enter_defer)
 
         # Start the background update.