fix known broken test

This commit is contained in:
Neil Johnson 2018-08-01 11:48:37 +01:00
parent 0aba3d361a
commit 4e6e00152c

View File

@ -15,7 +15,7 @@
from twisted.internet import defer from twisted.internet import defer
import tests.unittest
import tests.utils import tests.utils
@ -33,39 +33,34 @@ class InitTestCase(tests.unittest.TestCase):
self.store = hs.get_datastore() self.store = hs.get_datastore()
self.clock = hs.get_clock() self.clock = hs.get_clock()
@defer.inlineCallbacks
def test_count_monthly_users(self): def test_count_monthly_users(self):
count = self.store.count_monthly_users() count = yield self.store.count_monthly_users()
self.assertEqual(0, count) self.assertEqual(0, count)
self._insert_user_ips("@user:server1") yield self._insert_user_ips("@user:server1")
self._insert_user_ips("@user:server2") yield self._insert_user_ips("@user:server2")
count = self.store.count_monthly_users() count = yield self.store.count_monthly_users()
self.assertEqual(2, count) self.assertEqual(2, count)
@defer.inlineCallbacks
def _insert_user_ips(self, user): def _insert_user_ips(self, user):
""" """
Helper function to populate user_ips without using batch insertion infra Helper function to populate user_ips without using batch insertion infra
args: args:
user (str): specify username i.e. @user:server.com user (str): specify username i.e. @user:server.com
""" """
try: yield self.store._simple_upsert(
txn = self.store.db_conn.cursor() table="user_ips",
self.store.database_engine.lock_table(txn, "user_ips") keyvalues={
self.store._simple_upsert_txn( "user_id": user,
txn, "access_token": "access_token",
table="user_ips", "ip": "ip",
keyvalues={ "user_agent": "user_agent",
"user_id": user, "device_id": "device_id",
"access_token": "access_token", },
"ip": "ip", values={
"user_agent": "user_agent", "last_seen": self.clock.time_msec(),
"device_id": "device_id", }
}, )
values={
"last_seen": self.clock.time_msec(),
},
lock=False,
)
finally:
txn.close()