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.
|
