Fix unit tests

3 files changed, 15 insertions, 7 deletions

diff --git a/synapse/handlers/e2e_keys.py b/synapse/handlers/e2e_keys.py
index 49b277a1af..e40495d1ab 100644
--- a/synapse/handlers/e2e_keys.py
+++ b/synapse/handlers/e2e_keys.py
@@ -194,7 +194,7 @@ class E2eKeysHandler(object):
         # "unsigned" section
         for user_id, device_keys in results.items():
             for device_id, device_info in device_keys.items():
-                r = json.loads(device_info["key_json"])
+                r = dict(device_info["keys"])
                 r["unsigned"] = {}
                 display_name = device_info["device_display_name"]
                 if display_name is not None:
diff --git a/synapse/storage/end_to_end_keys.py b/synapse/storage/end_to_end_keys.py
index aa54d7637c..2040e022fa 100644
--- a/synapse/storage/end_to_end_keys.py
+++ b/synapse/storage/end_to_end_keys.py
@@ -15,6 +15,7 @@
 from twisted.internet import defer
 
 from canonicaljson import encode_canonical_json
+import ujson as json
 
 from ._base import SQLBaseStore
 
@@ -59,6 +60,7 @@ class EndToEndKeyStore(SQLBaseStore):
             "set_e2e_device_keys", _set_e2e_device_keys_txn
         )
 
+    @defer.inlineCallbacks
     def get_e2e_device_keys(self, query_list, include_all_devices=False):
         """Fetch a list of device keys.
         Args:
@@ -70,13 +72,19 @@ class EndToEndKeyStore(SQLBaseStore):
             dict containing "key_json", "device_display_name".
         """
         if not query_list:
-            return {}
+            defer.returnValue({})
 
-        return self.runInteraction(
+        results = yield self.runInteraction(
             "get_e2e_device_keys", self._get_e2e_device_keys_txn,
             query_list, include_all_devices,
         )
 
+        for user_id, device_keys in results.iteritems():
+            for device_id, device_info in device_keys.iteritems():
+                device_info["keys"] = json.loads(device_info.pop("key_json"))
+
+        defer.returnValue(results)
+
     def _get_e2e_device_keys_txn(self, txn, query_list, include_all_devices):
         query_clauses = []
         query_params = []
diff --git a/tests/storage/test_end_to_end_keys.py b/tests/storage/test_end_to_end_keys.py
index bfa6294250..84ce492a2c 100644
--- a/tests/storage/test_end_to_end_keys.py
+++ b/tests/storage/test_end_to_end_keys.py
@@ -33,7 +33,7 @@ class EndToEndKeyStoreTestCase(tests.unittest.TestCase):
     @defer.inlineCallbacks
     def test_key_without_device_name(self):
         now = 1470174257070
-        json = '{ "key": "value" }'
+        json = {"key": "value"}
 
         yield self.store.store_device(
             "user", "device", None
@@ -47,14 +47,14 @@ class EndToEndKeyStoreTestCase(tests.unittest.TestCase):
         self.assertIn("device", res["user"])
         dev = res["user"]["device"]
         self.assertDictContainsSubset({
-            "key_json": json,
+            "keys": json,
             "device_display_name": None,
         }, dev)
 
     @defer.inlineCallbacks
     def test_get_key_with_device_name(self):
         now = 1470174257070
-        json = '{ "key": "value" }'
+        json = {"key": "value"}
 
         yield self.store.set_e2e_device_keys(
             "user", "device", now, json)
@@ -67,7 +67,7 @@ class EndToEndKeyStoreTestCase(tests.unittest.TestCase):
         self.assertIn("device", res["user"])
         dev = res["user"]["device"]
         self.assertDictContainsSubset({
-            "key_json": json,
+            "keys": json,
             "device_display_name": "display_name",
         }, dev)