diff --git a/tests/storage/test_base.py b/tests/storage/test_base.py
index bf22540d99..64abe8cc49 100644
--- a/tests/storage/test_base.py
+++ b/tests/storage/test_base.py
@@ -148,8 +148,10 @@ class SQLBaseStoreTestCase(unittest.TestCase):
self.mock_txn.__iter__ = Mock(return_value=iter([(1,), (2,), (3,)]))
self.mock_txn.description = (("colA", None, None, None, None, None, None),)
- ret = yield self.datastore.db_pool.simple_select_list(
- table="tablename", keyvalues={"keycol": "A set"}, retcols=["colA"]
+ ret = yield defer.ensureDeferred(
+ self.datastore.db_pool.simple_select_list(
+ table="tablename", keyvalues={"keycol": "A set"}, retcols=["colA"]
+ )
)
self.assertEquals([{"colA": 1}, {"colA": 2}, {"colA": 3}], ret)
@@ -161,10 +163,12 @@ class SQLBaseStoreTestCase(unittest.TestCase):
def test_update_one_1col(self):
self.mock_txn.rowcount = 1
- yield self.datastore.db_pool.simple_update_one(
- table="tablename",
- keyvalues={"keycol": "TheKey"},
- updatevalues={"columnname": "New Value"},
+ yield defer.ensureDeferred(
+ self.datastore.db_pool.simple_update_one(
+ table="tablename",
+ keyvalues={"keycol": "TheKey"},
+ updatevalues={"columnname": "New Value"},
+ )
)
self.mock_txn.execute.assert_called_with(
@@ -176,10 +180,12 @@ class SQLBaseStoreTestCase(unittest.TestCase):
def test_update_one_4cols(self):
self.mock_txn.rowcount = 1
- yield self.datastore.db_pool.simple_update_one(
- table="tablename",
- keyvalues=OrderedDict([("colA", 1), ("colB", 2)]),
- updatevalues=OrderedDict([("colC", 3), ("colD", 4)]),
+ yield defer.ensureDeferred(
+ self.datastore.db_pool.simple_update_one(
+ table="tablename",
+ keyvalues=OrderedDict([("colA", 1), ("colB", 2)]),
+ updatevalues=OrderedDict([("colC", 3), ("colD", 4)]),
+ )
)
self.mock_txn.execute.assert_called_with(
diff --git a/tests/storage/test_directory.py b/tests/storage/test_directory.py
index daac947cb2..da93ca3980 100644
--- a/tests/storage/test_directory.py
+++ b/tests/storage/test_directory.py
@@ -42,7 +42,11 @@ class DirectoryStoreTestCase(unittest.TestCase):
self.assertEquals(
["#my-room:test"],
- (yield self.store.get_aliases_for_room(self.room.to_string())),
+ (
+ yield defer.ensureDeferred(
+ self.store.get_aliases_for_room(self.room.to_string())
+ )
+ ),
)
@defer.inlineCallbacks
diff --git a/tests/storage/test_main.py b/tests/storage/test_main.py
index fbf8af940a..954338a592 100644
--- a/tests/storage/test_main.py
+++ b/tests/storage/test_main.py
@@ -36,7 +36,9 @@ class DataStoreTestCase(unittest.TestCase):
def test_get_users_paginate(self):
yield self.store.register_user(self.user.to_string(), "pass")
yield defer.ensureDeferred(self.store.create_profile(self.user.localpart))
- yield self.store.set_profile_displayname(self.user.localpart, self.displayname)
+ yield defer.ensureDeferred(
+ self.store.set_profile_displayname(self.user.localpart, self.displayname)
+ )
users, total = yield self.store.get_users_paginate(
0, 10, name="bc", guests=False
|
