diff --git a/tests/storage/test__init__.py b/tests/storage/test__init__.py
index b763d395cc..c9ae349871 100644
--- a/tests/storage/test__init__.py
+++ b/tests/storage/test__init__.py
@@ -28,21 +28,45 @@ class InitTestCase(tests.unittest.TestCase):
@defer.inlineCallbacks
def setUp(self):
hs = yield tests.utils.setup_test_homeserver()
+
hs.config.max_mau_value = 50
hs.config.limit_usage_by_mau = True
self.store = hs.get_datastore()
+ self.clock = hs.get_clock()
- @defer.inlineCallbacks
def test_count_monthly_users(self):
- count = yield self.store.count_monthly_users()
+ count = self.store.count_monthly_users()
self.assertEqual(0, count)
- yield self.store.insert_client_ip(
- "@user:server1", "access_token", "ip", "user_agent", "device_id"
- )
- yield self.store.insert_client_ip(
- "@user:server2", "access_token", "ip", "user_agent", "device_id"
- )
- count = self.store.count_monthly_users()
+ self._insert_user_ips("@user:server1")
+ self._insert_user_ips("@user:server2")
+ count = self.store.count_monthly_users()
self.assertEqual(2, count)
+
+ def _insert_user_ips(self, user):
+ """
+ Helper function to populate user_ips without using batch insertion infra
+ args:
+ user (str): specify username i.e. @user:server.com
+ """
+ try:
+ txn = self.store.db_conn.cursor()
+ self.store.database_engine.lock_table(txn, "user_ips")
+ self.store._simple_upsert_txn(
+ txn,
+ table="user_ips",
+ keyvalues={
+ "user_id": user,
+ "access_token": "access_token",
+ "ip": "ip",
+ "user_agent": "user_agent",
+ "device_id": "device_id",
+ },
+ values={
+ "last_seen": self.clock.time_msec(),
+ },
+ lock=False,
+ )
+ finally:
+ txn.close()
|
