1 files changed, 21 insertions, 43 deletions
diff --git a/tests/storage/test_end_to_end_keys.py b/tests/storage/test_end_to_end_keys.py
index 84ce492a2c..8f0aaece40 100644
--- a/tests/storage/test_end_to_end_keys.py
+++ b/tests/storage/test_end_to_end_keys.py
@@ -26,8 +26,7 @@ class EndToEndKeyStoreTestCase(tests.unittest.TestCase):
 
     @defer.inlineCallbacks
     def setUp(self):
-        hs = yield tests.utils.setup_test_homeserver()
-
+        hs = yield tests.utils.setup_test_homeserver(self.addCleanup)
         self.store = hs.get_datastore()
 
     @defer.inlineCallbacks
@@ -35,70 +34,49 @@ class EndToEndKeyStoreTestCase(tests.unittest.TestCase):
         now = 1470174257070
         json = {"key": "value"}
 
-        yield self.store.store_device(
-            "user", "device", None
-        )
+        yield self.store.store_device("user", "device", None)
 
-        yield self.store.set_e2e_device_keys(
-            "user", "device", now, json)
+        yield self.store.set_e2e_device_keys("user", "device", now, json)
 
         res = yield self.store.get_e2e_device_keys((("user", "device"),))
         self.assertIn("user", res)
         self.assertIn("device", res["user"])
         dev = res["user"]["device"]
-        self.assertDictContainsSubset({
-            "keys": json,
-            "device_display_name": None,
-        }, dev)
+        self.assertDictContainsSubset({"keys": json, "device_display_name": None}, dev)
 
     @defer.inlineCallbacks
     def test_get_key_with_device_name(self):
         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 self.store.set_e2e_device_keys("user", "device", now, json)
+        yield self.store.store_device("user", "device", "display_name")
 
         res = yield self.store.get_e2e_device_keys((("user", "device"),))
         self.assertIn("user", res)
         self.assertIn("device", res["user"])
         dev = res["user"]["device"]
-        self.assertDictContainsSubset({
-            "keys": json,
-            "device_display_name": "display_name",
-        }, dev)
+        self.assertDictContainsSubset(
+            {"keys": json, "device_display_name": "display_name"}, dev
+        )
 
     @defer.inlineCallbacks
     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 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 self.store.set_e2e_device_keys(
-            "user1", "device1", now, 'json11')
-        yield self.store.set_e2e_device_keys(
-            "user1", "device2", now, 'json12')
-        yield self.store.set_e2e_device_keys(
-            "user2", "device1", now, 'json21')
-        yield self.store.set_e2e_device_keys(
-            "user2", "device2", now, 'json22')
-
-        res = yield self.store.get_e2e_device_keys((("user1", "device1"),
-                                                    ("user2", "device2")))
+        yield self.store.set_e2e_device_keys("user1", "device1", now, 'json11')
+        yield self.store.set_e2e_device_keys("user1", "device2", now, 'json12')
+        yield self.store.set_e2e_device_keys("user2", "device1", now, 'json21')
+        yield self.store.set_e2e_device_keys("user2", "device2", now, 'json22')
+
+        res = yield self.store.get_e2e_device_keys(
+            (("user1", "device1"), ("user2", "device2"))
+        )
         self.assertIn("user1", res)
         self.assertIn("device1", res["user1"])
         self.assertNotIn("device2", res["user1"])