fix user_ips counting

1 files changed, 33 insertions, 9 deletions

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()