diff --git a/tests/storage/test_devices.py b/tests/storage/test_devices.py
index c2539b353a..ecb00f4e02 100644
--- a/tests/storage/test_devices.py
+++ b/tests/storage/test_devices.py
@@ -23,7 +23,7 @@ import tests.utils
class DeviceStoreTestCase(tests.unittest.TestCase):
def __init__(self, *args, **kwargs):
- super(DeviceStoreTestCase, self).__init__(*args, **kwargs)
+ super().__init__(*args, **kwargs)
self.store = None # type: synapse.storage.DataStore
@defer.inlineCallbacks
@@ -34,9 +34,11 @@ class DeviceStoreTestCase(tests.unittest.TestCase):
@defer.inlineCallbacks
def test_store_new_device(self):
- yield self.store.store_device("user_id", "device_id", "display_name")
+ yield defer.ensureDeferred(
+ self.store.store_device("user_id", "device_id", "display_name")
+ )
- res = yield self.store.get_device("user_id", "device_id")
+ res = yield defer.ensureDeferred(self.store.get_device("user_id", "device_id"))
self.assertDictContainsSubset(
{
"user_id": "user_id",
@@ -48,11 +50,17 @@ class DeviceStoreTestCase(tests.unittest.TestCase):
@defer.inlineCallbacks
def test_get_devices_by_user(self):
- yield self.store.store_device("user_id", "device1", "display_name 1")
- yield self.store.store_device("user_id", "device2", "display_name 2")
- yield self.store.store_device("user_id2", "device3", "display_name 3")
+ yield defer.ensureDeferred(
+ self.store.store_device("user_id", "device1", "display_name 1")
+ )
+ yield defer.ensureDeferred(
+ self.store.store_device("user_id", "device2", "display_name 2")
+ )
+ yield defer.ensureDeferred(
+ self.store.store_device("user_id2", "device3", "display_name 3")
+ )
- res = yield self.store.get_devices_by_user("user_id")
+ res = yield defer.ensureDeferred(self.store.get_devices_by_user("user_id"))
self.assertEqual(2, len(res.keys()))
self.assertDictContainsSubset(
{
@@ -76,13 +84,13 @@ class DeviceStoreTestCase(tests.unittest.TestCase):
device_ids = ["device_id1", "device_id2"]
# Add two device updates with a single stream_id
- yield self.store.add_device_change_to_streams(
- "user_id", device_ids, ["somehost"]
+ yield defer.ensureDeferred(
+ self.store.add_device_change_to_streams("user_id", device_ids, ["somehost"])
)
# Get all device updates ever meant for this remote
- now_stream_id, device_updates = yield self.store.get_device_updates_by_remote(
- "somehost", -1, limit=100
+ now_stream_id, device_updates = yield defer.ensureDeferred(
+ self.store.get_device_updates_by_remote("somehost", -1, limit=100)
)
# Check original device_ids are contained within these updates
@@ -99,29 +107,35 @@ class DeviceStoreTestCase(tests.unittest.TestCase):
@defer.inlineCallbacks
def test_update_device(self):
- yield self.store.store_device("user_id", "device_id", "display_name 1")
+ yield defer.ensureDeferred(
+ self.store.store_device("user_id", "device_id", "display_name 1")
+ )
- res = yield self.store.get_device("user_id", "device_id")
+ res = yield defer.ensureDeferred(self.store.get_device("user_id", "device_id"))
self.assertEqual("display_name 1", res["display_name"])
# do a no-op first
- yield self.store.update_device("user_id", "device_id")
- res = yield self.store.get_device("user_id", "device_id")
+ yield defer.ensureDeferred(self.store.update_device("user_id", "device_id"))
+ res = yield defer.ensureDeferred(self.store.get_device("user_id", "device_id"))
self.assertEqual("display_name 1", res["display_name"])
# do the update
- yield self.store.update_device(
- "user_id", "device_id", new_display_name="display_name 2"
+ yield defer.ensureDeferred(
+ self.store.update_device(
+ "user_id", "device_id", new_display_name="display_name 2"
+ )
)
# check it worked
- res = yield self.store.get_device("user_id", "device_id")
+ res = yield defer.ensureDeferred(self.store.get_device("user_id", "device_id"))
self.assertEqual("display_name 2", res["display_name"])
@defer.inlineCallbacks
def test_update_unknown_device(self):
with self.assertRaises(synapse.api.errors.StoreError) as cm:
- yield self.store.update_device(
- "user_id", "unknown_device_id", new_display_name="display_name 2"
+ yield defer.ensureDeferred(
+ self.store.update_device(
+ "user_id", "unknown_device_id", new_display_name="display_name 2"
+ )
)
self.assertEqual(404, cm.exception.code)
|
