Merge branch 'develop' into develop

1 files changed, 37 insertions, 67 deletions

diff --git a/tests/handlers/test_e2e_keys.py b/tests/handlers/test_e2e_keys.py
index d1bd87b898..8dccc6826e 100644
--- a/tests/handlers/test_e2e_keys.py
+++ b/tests/handlers/test_e2e_keys.py
@@ -14,27 +14,27 @@
 # limitations under the License.
 
 import mock
-from synapse.api import errors
+
 from twisted.internet import defer
 
 import synapse.api.errors
 import synapse.handlers.e2e_keys
-
 import synapse.storage
+from synapse.api import errors
+
 from tests import unittest, utils
 
 
 class E2eKeysHandlerTestCase(unittest.TestCase):
     def __init__(self, *args, **kwargs):
         super(E2eKeysHandlerTestCase, self).__init__(*args, **kwargs)
-        self.hs = None       # type: synapse.server.HomeServer
+        self.hs = None  # type: synapse.server.HomeServer
         self.handler = None  # type: synapse.handlers.e2e_keys.E2eKeysHandler
 
     @defer.inlineCallbacks
     def setUp(self):
         self.hs = yield utils.setup_test_homeserver(
-            handlers=None,
-            federation_client=mock.Mock(),
+            self.addCleanup, handlers=None, federation_client=mock.Mock()
         )
         self.handler = synapse.handlers.e2e_keys.E2eKeysHandler(self.hs)
 
@@ -53,30 +53,21 @@ class E2eKeysHandlerTestCase(unittest.TestCase):
         device_id = "xyz"
         keys = {
             "alg1:k1": "key1",
-            "alg2:k2": {
-                "key": "key2",
-                "signatures": {"k1": "sig1"}
-            },
-            "alg2:k3": {
-                "key": "key3",
-            },
+            "alg2:k2": {"key": "key2", "signatures": {"k1": "sig1"}},
+            "alg2:k3": {"key": "key3"},
         }
 
         res = yield self.handler.upload_keys_for_user(
-            local_user, device_id, {"one_time_keys": keys},
+            local_user, device_id, {"one_time_keys": keys}
         )
-        self.assertDictEqual(res, {
-            "one_time_key_counts": {"alg1": 1, "alg2": 2}
-        })
+        self.assertDictEqual(res, {"one_time_key_counts": {"alg1": 1, "alg2": 2}})
 
         # we should be able to change the signature without a problem
         keys["alg2:k2"]["signatures"]["k1"] = "sig2"
         res = yield self.handler.upload_keys_for_user(
-            local_user, device_id, {"one_time_keys": keys},
+            local_user, device_id, {"one_time_keys": keys}
         )
-        self.assertDictEqual(res, {
-            "one_time_key_counts": {"alg1": 1, "alg2": 2}
-        })
+        self.assertDictEqual(res, {"one_time_key_counts": {"alg1": 1, "alg2": 2}})
 
     @defer.inlineCallbacks
     def test_change_one_time_keys(self):
@@ -86,25 +77,18 @@ class E2eKeysHandlerTestCase(unittest.TestCase):
         device_id = "xyz"
         keys = {
             "alg1:k1": "key1",
-            "alg2:k2": {
-                "key": "key2",
-                "signatures": {"k1": "sig1"}
-            },
-            "alg2:k3": {
-                "key": "key3",
-            },
+            "alg2:k2": {"key": "key2", "signatures": {"k1": "sig1"}},
+            "alg2:k3": {"key": "key3"},
         }
 
         res = yield self.handler.upload_keys_for_user(
-            local_user, device_id, {"one_time_keys": keys},
+            local_user, device_id, {"one_time_keys": keys}
         )
-        self.assertDictEqual(res, {
-            "one_time_key_counts": {"alg1": 1, "alg2": 2}
-        })
+        self.assertDictEqual(res, {"one_time_key_counts": {"alg1": 1, "alg2": 2}})
 
         try:
             yield self.handler.upload_keys_for_user(
-                local_user, device_id, {"one_time_keys": {"alg1:k1": "key2"}},
+                local_user, device_id, {"one_time_keys": {"alg1:k1": "key2"}}
             )
             self.fail("No error when changing string key")
         except errors.SynapseError:
@@ -112,7 +96,7 @@ class E2eKeysHandlerTestCase(unittest.TestCase):
 
         try:
             yield self.handler.upload_keys_for_user(
-                local_user, device_id, {"one_time_keys": {"alg2:k3": "key2"}},
+                local_user, device_id, {"one_time_keys": {"alg2:k3": "key2"}}
             )
             self.fail("No error when replacing dict key with string")
         except errors.SynapseError:
@@ -120,9 +104,7 @@ class E2eKeysHandlerTestCase(unittest.TestCase):
 
         try:
             yield self.handler.upload_keys_for_user(
-                local_user, device_id, {
-                    "one_time_keys": {"alg1:k1": {"key": "key"}}
-                },
+                local_user, device_id, {"one_time_keys": {"alg1:k1": {"key": "key"}}}
             )
             self.fail("No error when replacing string key with dict")
         except errors.SynapseError:
@@ -130,13 +112,12 @@ class E2eKeysHandlerTestCase(unittest.TestCase):
 
         try:
             yield self.handler.upload_keys_for_user(
-                local_user, device_id, {
+                local_user,
+                device_id,
+                {
                     "one_time_keys": {
-                        "alg2:k2": {
-                            "key": "key3",
-                            "signatures": {"k1": "sig1"},
-                        }
-                    },
+                        "alg2:k2": {"key": "key3", "signatures": {"k1": "sig1"}}
+                    }
                 },
             )
             self.fail("No error when replacing dict key")
@@ -147,31 +128,20 @@ class E2eKeysHandlerTestCase(unittest.TestCase):
     def test_claim_one_time_key(self):
         local_user = "@boris:" + self.hs.hostname
         device_id = "xyz"
-        keys = {
-            "alg1:k1": "key1",
-        }
+        keys = {"alg1:k1": "key1"}
 
         res = yield self.handler.upload_keys_for_user(
-            local_user, device_id, {"one_time_keys": keys},
+            local_user, device_id, {"one_time_keys": keys}
+        )
+        self.assertDictEqual(res, {"one_time_key_counts": {"alg1": 1}})
+
+        res2 = yield self.handler.claim_one_time_keys(
+            {"one_time_keys": {local_user: {device_id: "alg1"}}}, timeout=None
+        )
+        self.assertEqual(
+            res2,
+            {
+                "failures": {},
+                "one_time_keys": {local_user: {device_id: {"alg1:k1": "key1"}}},
+            },
         )
-        self.assertDictEqual(res, {
-            "one_time_key_counts": {"alg1": 1}
-        })
-
-        res2 = yield self.handler.claim_one_time_keys({
-            "one_time_keys": {
-                local_user: {
-                    device_id: "alg1"
-                }
-            }
-        }, timeout=None)
-        self.assertEqual(res2, {
-            "failures": {},
-            "one_time_keys": {
-                local_user: {
-                    device_id: {
-                        "alg1:k1": "key1"
-                    }
-                }
-            }
-        })