diff --git a/tests/handlers/test_stats.py b/tests/handlers/test_stats.py
index 7569b6fab5..d9d312f0fb 100644
--- a/tests/handlers/test_stats.py
+++ b/tests/handlers/test_stats.py
@@ -13,9 +13,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from synapse import storage
from synapse.rest import admin
from synapse.rest.client.v1 import login, room
+from synapse.storage.data_stores.main import stats
from tests import unittest
@@ -42,16 +42,16 @@ 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._simple_insert(
+ self.store.db.simple_insert(
"background_updates",
{"update_name": "populate_stats_prepare", "progress_json": "{}"},
)
)
self.get_success(
- self.store._simple_insert(
+ self.store.db.simple_insert(
"background_updates",
{
"update_name": "populate_stats_process_rooms",
@@ -61,7 +61,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
)
)
self.get_success(
- self.store._simple_insert(
+ self.store.db.simple_insert(
"background_updates",
{
"update_name": "populate_stats_process_users",
@@ -71,7 +71,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
)
)
self.get_success(
- self.store._simple_insert(
+ self.store.db.simple_insert(
"background_updates",
{
"update_name": "populate_stats_cleanup",
@@ -82,21 +82,21 @@ class StatsRoomTests(unittest.HomeserverTestCase):
)
def get_all_room_state(self):
- return self.store._simple_select_list(
+ return self.store.db.simple_select_list(
"room_stats_state", None, retcols=("name", "topic", "canonical_alias")
)
def _get_current_stats(self, stats_type, stat_id):
- table, id_col = storage.stats.TYPE_TO_TABLE[stats_type]
+ table, id_col = stats.TYPE_TO_TABLE[stats_type]
- cols = list(storage.stats.ABSOLUTE_STATS_FIELDS[stats_type]) + list(
- storage.stats.PER_SLICE_FIELDS[stats_type]
+ cols = list(stats.ABSOLUTE_STATS_FIELDS[stats_type]) + list(
+ stats.PER_SLICE_FIELDS[stats_type]
)
end_ts = self.store.quantise_stats_time(self.reactor.seconds() * 1000)
return self.get_success(
- self.store._simple_select_one(
+ self.store.db.simple_select_one(
table + "_historical",
{id_col: stat_id, end_ts: end_ts},
cols,
@@ -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,9 +186,9 @@ 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._simple_update_one(
+ self.store.db.simple_update_one(
table="stats_incremental_position",
keyvalues={},
updatevalues={"stream_id": 0},
@@ -188,14 +196,18 @@ class StatsRoomTests(unittest.HomeserverTestCase):
)
self.get_success(
- self.store._simple_insert(
+ self.store.db.simple_insert(
"background_updates",
{"update_name": "populate_stats_prepare", "progress_json": "{}"},
)
)
- 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)
@@ -205,13 +217,13 @@ class StatsRoomTests(unittest.HomeserverTestCase):
# Now do the initial ingestion.
self.get_success(
- self.store._simple_insert(
+ self.store.db.simple_insert(
"background_updates",
{"update_name": "populate_stats_process_rooms", "progress_json": "{}"},
)
)
self.get_success(
- self.store._simple_insert(
+ self.store.db.simple_insert(
"background_updates",
{
"update_name": "populate_stats_cleanup",
@@ -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)
@@ -607,6 +623,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
"""
self.hs.config.stats_enabled = False
+ self.handler.stats_enabled = False
u1 = self.register_user("u1", "pass")
u1token = self.login("u1", "pass")
@@ -618,6 +635,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
self.assertIsNone(self._get_current_stats("user", u1))
self.hs.config.stats_enabled = True
+ self.handler.stats_enabled = True
self._perform_background_initial_update()
@@ -651,15 +669,15 @@ 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._simple_delete(
+ self.store.db.simple_delete(
"room_stats_current", {"1": 1}, "test_delete_stats"
)
)
self.get_success(
- self.store._simple_delete(
+ self.store.db.simple_delete(
"user_stats_current", {"1": 1}, "test_delete_stats"
)
)
@@ -671,9 +689,9 @@ 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._simple_insert(
+ self.store.db.simple_insert(
"background_updates",
{
"update_name": "populate_stats_process_rooms",
@@ -683,7 +701,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
)
)
self.get_success(
- self.store._simple_insert(
+ self.store.db.simple_insert(
"background_updates",
{
"update_name": "populate_stats_process_users",
@@ -693,7 +711,7 @@ class StatsRoomTests(unittest.HomeserverTestCase):
)
)
self.get_success(
- self.store._simple_insert(
+ self.store.db.simple_insert(
"background_updates",
{
"update_name": "populate_stats_cleanup",
@@ -703,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)
|
