diff --git a/tests/handlers/test_stats.py b/tests/handlers/test_stats.py
index 7f7962c3dd..d9d312f0fb 100644
--- a/tests/handlers/test_stats.py
+++ b/tests/handlers/test_stats.py
@@ -42,7 +42,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
Add the background updates we need to run.
"""
# Ugh, have to reset this flag
- self.store._all_done = False
+ self.store.db.updates._all_done = False
self.get_success(
self.store.db.simple_insert(
@@ -108,8 +108,12 @@ class StatsRoomTests(unittest.HomeserverTestCase):
# Do the initial population of the stats via the background update
self._add_background_updates()
- while not self.get_success(self.store.has_completed_background_updates()):
- self.get_success(self.store.do_next_background_update(100), by=0.1)
+ while not self.get_success(
+ self.store.db.updates.has_completed_background_updates()
+ ):
+ self.get_success(
+ self.store.db.updates.do_next_background_update(100), by=0.1
+ )
def test_initial_room(self):
"""
@@ -141,8 +145,12 @@ class StatsRoomTests(unittest.HomeserverTestCase):
# Do the initial population of the user directory via the background update
self._add_background_updates()
- while not self.get_success(self.store.has_completed_background_updates()):
- self.get_success(self.store.do_next_background_update(100), by=0.1)
+ while not self.get_success(
+ self.store.db.updates.has_completed_background_updates()
+ ):
+ self.get_success(
+ self.store.db.updates.do_next_background_update(100), by=0.1
+ )
r = self.get_success(self.get_all_room_state())
@@ -178,7 +186,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
# the position that the deltas should begin at, once they take over.
self.hs.config.stats_enabled = True
self.handler.stats_enabled = True
- self.store._all_done = False
+ self.store.db.updates._all_done = False
self.get_success(
self.store.db.simple_update_one(
table="stats_incremental_position",
@@ -194,8 +202,12 @@ class StatsRoomTests(unittest.HomeserverTestCase):
)
)
- while not self.get_success(self.store.has_completed_background_updates()):
- self.get_success(self.store.do_next_background_update(100), by=0.1)
+ while not self.get_success(
+ self.store.db.updates.has_completed_background_updates()
+ ):
+ self.get_success(
+ self.store.db.updates.do_next_background_update(100), by=0.1
+ )
# Now, before the table is actually ingested, add some more events.
self.helper.invite(room=room_1, src=u1, targ=u2, tok=u1_token)
@@ -221,9 +233,13 @@ class StatsRoomTests(unittest.HomeserverTestCase):
)
)
- self.store._all_done = False
- while not self.get_success(self.store.has_completed_background_updates()):
- self.get_success(self.store.do_next_background_update(100), by=0.1)
+ self.store.db.updates._all_done = False
+ while not self.get_success(
+ self.store.db.updates.has_completed_background_updates()
+ ):
+ self.get_success(
+ self.store.db.updates.do_next_background_update(100), by=0.1
+ )
self.reactor.advance(86401)
@@ -653,7 +669,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
# preparation stage of the initial background update
# Ugh, have to reset this flag
- self.store._all_done = False
+ self.store.db.updates._all_done = False
self.get_success(
self.store.db.simple_delete(
@@ -673,7 +689,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
# now do the background updates
- self.store._all_done = False
+ self.store.db.updates._all_done = False
self.get_success(
self.store.db.simple_insert(
"background_updates",
@@ -705,8 +721,12 @@ class StatsRoomTests(unittest.HomeserverTestCase):
)
)
- while not self.get_success(self.store.has_completed_background_updates()):
- self.get_success(self.store.do_next_background_update(100), by=0.1)
+ while not self.get_success(
+ self.store.db.updates.has_completed_background_updates()
+ ):
+ self.get_success(
+ self.store.db.updates.do_next_background_update(100), by=0.1
+ )
r1stats_complete = self._get_current_stats("room", r1)
u1stats_complete = self._get_current_stats("user", u1)
diff --git a/tests/handlers/test_user_directory.py b/tests/handlers/test_user_directory.py
index bc9d441541..26071059d2 100644
--- a/tests/handlers/test_user_directory.py
+++ b/tests/handlers/test_user_directory.py
@@ -181,7 +181,7 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
Add the background updates we need to run.
"""
# Ugh, have to reset this flag
- self.store._all_done = False
+ self.store.db.updates._all_done = False
self.get_success(
self.store.db.simple_insert(
@@ -255,8 +255,12 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
# Do the initial population of the user directory via the background update
self._add_background_updates()
- while not self.get_success(self.store.has_completed_background_updates()):
- self.get_success(self.store.do_next_background_update(100), by=0.1)
+ while not self.get_success(
+ self.store.db.updates.has_completed_background_updates()
+ ):
+ self.get_success(
+ self.store.db.updates.do_next_background_update(100), by=0.1
+ )
shares_private = self.get_users_who_share_private_rooms()
public_users = self.get_users_in_public_rooms()
@@ -290,8 +294,12 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
# Do the initial population of the user directory via the background update
self._add_background_updates()
- while not self.get_success(self.store.has_completed_background_updates()):
- self.get_success(self.store.do_next_background_update(100), by=0.1)
+ while not self.get_success(
+ self.store.db.updates.has_completed_background_updates()
+ ):
+ self.get_success(
+ self.store.db.updates.do_next_background_update(100), by=0.1
+ )
shares_private = self.get_users_who_share_private_rooms()
public_users = self.get_users_in_public_rooms()
diff --git a/tests/storage/test_background_update.py b/tests/storage/test_background_update.py
index e360297df9..aec76f4ab1 100644
--- a/tests/storage/test_background_update.py
+++ b/tests/storage/test_background_update.py
@@ -15,7 +15,7 @@ class BackgroundUpdateTestCase(unittest.TestCase):
self.update_handler = Mock()
- yield self.store.register_background_update_handler(
+ yield self.store.db.updates.register_background_update_handler(
"test_update", self.update_handler
)
@@ -23,7 +23,7 @@ class BackgroundUpdateTestCase(unittest.TestCase):
# (perhaps we should run them as part of the test HS setup, since we
# run all of the other schema setup stuff there?)
while True:
- res = yield self.store.do_next_background_update(1000)
+ res = yield self.store.db.updates.do_next_background_update(1000)
if res is None:
break
@@ -39,7 +39,7 @@ class BackgroundUpdateTestCase(unittest.TestCase):
progress = {"my_key": progress["my_key"] + 1}
yield self.store.db.runInteraction(
"update_progress",
- self.store._background_update_progress_txn,
+ self.store.db.updates._background_update_progress_txn,
"test_update",
progress,
)
@@ -47,29 +47,37 @@ class BackgroundUpdateTestCase(unittest.TestCase):
self.update_handler.side_effect = update
- yield self.store.start_background_update("test_update", {"my_key": 1})
+ yield self.store.db.updates.start_background_update(
+ "test_update", {"my_key": 1}
+ )
self.update_handler.reset_mock()
- result = yield self.store.do_next_background_update(duration_ms * desired_count)
+ result = yield self.store.db.updates.do_next_background_update(
+ duration_ms * desired_count
+ )
self.assertIsNotNone(result)
self.update_handler.assert_called_once_with(
- {"my_key": 1}, self.store.DEFAULT_BACKGROUND_BATCH_SIZE
+ {"my_key": 1}, self.store.db.updates.DEFAULT_BACKGROUND_BATCH_SIZE
)
# second step: complete the update
@defer.inlineCallbacks
def update(progress, count):
- yield self.store._end_background_update("test_update")
+ yield self.store.db.updates._end_background_update("test_update")
return count
self.update_handler.side_effect = update
self.update_handler.reset_mock()
- result = yield self.store.do_next_background_update(duration_ms * desired_count)
+ result = yield self.store.db.updates.do_next_background_update(
+ duration_ms * desired_count
+ )
self.assertIsNotNone(result)
self.update_handler.assert_called_once_with({"my_key": 2}, desired_count)
# third step: we don't expect to be called any more
self.update_handler.reset_mock()
- result = yield self.store.do_next_background_update(duration_ms * desired_count)
+ result = yield self.store.db.updates.do_next_background_update(
+ duration_ms * desired_count
+ )
self.assertIsNone(result)
self.assertFalse(self.update_handler.called)
diff --git a/tests/storage/test_cleanup_extrems.py b/tests/storage/test_cleanup_extrems.py
index e454bbff29..029ac26454 100644
--- a/tests/storage/test_cleanup_extrems.py
+++ b/tests/storage/test_cleanup_extrems.py
@@ -46,7 +46,9 @@ class CleanupExtremBackgroundUpdateStoreTestCase(HomeserverTestCase):
"""Re run the background update to clean up the extremities.
"""
# Make sure we don't clash with in progress updates.
- self.assertTrue(self.store._all_done, "Background updates are still ongoing")
+ self.assertTrue(
+ self.store.db.updates._all_done, "Background updates are still ongoing"
+ )
schema_path = os.path.join(
prepare_database.dir_path,
@@ -68,10 +70,14 @@ class CleanupExtremBackgroundUpdateStoreTestCase(HomeserverTestCase):
)
# Ugh, have to reset this flag
- self.store._all_done = False
+ self.store.db.updates._all_done = False
- while not self.get_success(self.store.has_completed_background_updates()):
- self.get_success(self.store.do_next_background_update(100), by=0.1)
+ while not self.get_success(
+ self.store.db.updates.has_completed_background_updates()
+ ):
+ self.get_success(
+ self.store.db.updates.do_next_background_update(100), by=0.1
+ )
def test_soft_failed_extremities_handled_correctly(self):
"""Test that extremities are correctly calculated in the presence of
diff --git a/tests/storage/test_client_ips.py b/tests/storage/test_client_ips.py
index c4f838907c..fc279340d4 100644
--- a/tests/storage/test_client_ips.py
+++ b/tests/storage/test_client_ips.py
@@ -202,8 +202,12 @@ class ClientIpStoreTestCase(unittest.HomeserverTestCase):
def test_devices_last_seen_bg_update(self):
# First make sure we have completed all updates.
- while not self.get_success(self.store.has_completed_background_updates()):
- self.get_success(self.store.do_next_background_update(100), by=0.1)
+ while not self.get_success(
+ self.store.db.updates.has_completed_background_updates()
+ ):
+ self.get_success(
+ self.store.db.updates.do_next_background_update(100), by=0.1
+ )
# Insert a user IP
user_id = "@user:id"
@@ -256,11 +260,15 @@ class ClientIpStoreTestCase(unittest.HomeserverTestCase):
)
# ... and tell the DataStore that it hasn't finished all updates yet
- self.store._all_done = False
+ self.store.db.updates._all_done = False
# Now let's actually drive the updates to completion
- while not self.get_success(self.store.has_completed_background_updates()):
- self.get_success(self.store.do_next_background_update(100), by=0.1)
+ while not self.get_success(
+ self.store.db.updates.has_completed_background_updates()
+ ):
+ self.get_success(
+ self.store.db.updates.do_next_background_update(100), by=0.1
+ )
# We should now get the correct result again
result = self.get_success(
@@ -281,8 +289,12 @@ class ClientIpStoreTestCase(unittest.HomeserverTestCase):
def test_old_user_ips_pruned(self):
# First make sure we have completed all updates.
- while not self.get_success(self.store.has_completed_background_updates()):
- self.get_success(self.store.do_next_background_update(100), by=0.1)
+ while not self.get_success(
+ self.store.db.updates.has_completed_background_updates()
+ ):
+ self.get_success(
+ self.store.db.updates.do_next_background_update(100), by=0.1
+ )
# Insert a user IP
user_id = "@user:id"
diff --git a/tests/storage/test_roommember.py b/tests/storage/test_roommember.py
index 5f957680a2..7840f63fe3 100644
--- a/tests/storage/test_roommember.py
+++ b/tests/storage/test_roommember.py
@@ -122,8 +122,12 @@ class CurrentStateMembershipUpdateTestCase(unittest.HomeserverTestCase):
def test_can_rerun_update(self):
# First make sure we have completed all updates.
- while not self.get_success(self.store.has_completed_background_updates()):
- self.get_success(self.store.do_next_background_update(100), by=0.1)
+ while not self.get_success(
+ self.store.db.updates.has_completed_background_updates()
+ ):
+ self.get_success(
+ self.store.db.updates.do_next_background_update(100), by=0.1
+ )
# Now let's create a room, which will insert a membership
user = UserID("alice", "test")
@@ -143,8 +147,12 @@ class CurrentStateMembershipUpdateTestCase(unittest.HomeserverTestCase):
)
# ... and tell the DataStore that it hasn't finished all updates yet
- self.store._all_done = False
+ self.store.db.updates._all_done = False
# Now let's actually drive the updates to completion
- while not self.get_success(self.store.has_completed_background_updates()):
- self.get_success(self.store.do_next_background_update(100), by=0.1)
+ while not self.get_success(
+ self.store.db.updates.has_completed_background_updates()
+ ):
+ self.get_success(
+ self.store.db.updates.do_next_background_update(100), by=0.1
+ )
diff --git a/tests/unittest.py b/tests/unittest.py
index fc856a574a..68d245ec9f 100644
--- a/tests/unittest.py
+++ b/tests/unittest.py
@@ -401,10 +401,12 @@ class HomeserverTestCase(TestCase):
hs = setup_test_homeserver(self.addCleanup, *args, **kwargs)
stor = hs.get_datastore()
- # Run the database background updates.
- if hasattr(stor, "do_next_background_update"):
- while not self.get_success(stor.has_completed_background_updates()):
- self.get_success(stor.do_next_background_update(1))
+ # Run the database background updates, when running against "master".
+ if hs.__class__.__name__ == "TestHomeServer":
+ while not self.get_success(
+ stor.db.updates.has_completed_background_updates()
+ ):
+ self.get_success(stor.db.updates.do_next_background_update(1))
return hs
|
