diff --git a/tests/storage/test_end_to_end_keys.py b/tests/storage/test_end_to_end_keys.py
index 398d546280..3fc4bb13b6 100644
--- a/tests/storage/test_end_to_end_keys.py
+++ b/tests/storage/test_end_to_end_keys.py
@@ -30,11 +30,15 @@ class EndToEndKeyStoreTestCase(tests.unittest.TestCase):
now = 1470174257070
json = {"key": "value"}
- yield self.store.store_device("user", "device", None)
+ yield defer.ensureDeferred(self.store.store_device("user", "device", None))
- yield self.store.set_e2e_device_keys("user", "device", now, json)
+ yield defer.ensureDeferred(
+ self.store.set_e2e_device_keys("user", "device", now, json)
+ )
- res = yield self.store.get_e2e_device_keys((("user", "device"),))
+ res = yield defer.ensureDeferred(
+ self.store.get_e2e_device_keys_for_cs_api((("user", "device"),))
+ )
self.assertIn("user", res)
self.assertIn("device", res["user"])
dev = res["user"]["device"]
@@ -45,14 +49,18 @@ class EndToEndKeyStoreTestCase(tests.unittest.TestCase):
now = 1470174257070
json = {"key": "value"}
- yield self.store.store_device("user", "device", None)
+ yield defer.ensureDeferred(self.store.store_device("user", "device", None))
- changed = yield self.store.set_e2e_device_keys("user", "device", now, json)
+ changed = yield defer.ensureDeferred(
+ self.store.set_e2e_device_keys("user", "device", now, json)
+ )
self.assertTrue(changed)
# If we try to upload the same key then we should be told nothing
# changed
- changed = yield self.store.set_e2e_device_keys("user", "device", now, json)
+ changed = yield defer.ensureDeferred(
+ self.store.set_e2e_device_keys("user", "device", now, json)
+ )
self.assertFalse(changed)
@defer.inlineCallbacks
@@ -60,10 +68,16 @@ class EndToEndKeyStoreTestCase(tests.unittest.TestCase):
now = 1470174257070
json = {"key": "value"}
- yield self.store.set_e2e_device_keys("user", "device", now, json)
- yield self.store.store_device("user", "device", "display_name")
+ yield defer.ensureDeferred(
+ self.store.set_e2e_device_keys("user", "device", now, json)
+ )
+ yield defer.ensureDeferred(
+ self.store.store_device("user", "device", "display_name")
+ )
- res = yield self.store.get_e2e_device_keys((("user", "device"),))
+ res = yield defer.ensureDeferred(
+ self.store.get_e2e_device_keys_for_cs_api((("user", "device"),))
+ )
self.assertIn("user", res)
self.assertIn("device", res["user"])
dev = res["user"]["device"]
@@ -75,18 +89,28 @@ class EndToEndKeyStoreTestCase(tests.unittest.TestCase):
def test_multiple_devices(self):
now = 1470174257070
- yield self.store.store_device("user1", "device1", None)
- yield self.store.store_device("user1", "device2", None)
- yield self.store.store_device("user2", "device1", None)
- yield self.store.store_device("user2", "device2", None)
+ yield defer.ensureDeferred(self.store.store_device("user1", "device1", None))
+ yield defer.ensureDeferred(self.store.store_device("user1", "device2", None))
+ yield defer.ensureDeferred(self.store.store_device("user2", "device1", None))
+ yield defer.ensureDeferred(self.store.store_device("user2", "device2", None))
- yield self.store.set_e2e_device_keys("user1", "device1", now, {"key": "json11"})
- yield self.store.set_e2e_device_keys("user1", "device2", now, {"key": "json12"})
- yield self.store.set_e2e_device_keys("user2", "device1", now, {"key": "json21"})
- yield self.store.set_e2e_device_keys("user2", "device2", now, {"key": "json22"})
+ yield defer.ensureDeferred(
+ self.store.set_e2e_device_keys("user1", "device1", now, {"key": "json11"})
+ )
+ yield defer.ensureDeferred(
+ self.store.set_e2e_device_keys("user1", "device2", now, {"key": "json12"})
+ )
+ yield defer.ensureDeferred(
+ self.store.set_e2e_device_keys("user2", "device1", now, {"key": "json21"})
+ )
+ yield defer.ensureDeferred(
+ self.store.set_e2e_device_keys("user2", "device2", now, {"key": "json22"})
+ )
- res = yield self.store.get_e2e_device_keys(
- (("user1", "device1"), ("user2", "device2"))
+ res = yield defer.ensureDeferred(
+ self.store.get_e2e_device_keys_for_cs_api(
+ (("user1", "device1"), ("user2", "device2"))
+ )
)
self.assertIn("user1", res)
self.assertIn("device1", res["user1"])
|
