Fix unit tests

This commit is contained in:
Erik Johnston 2017-01-30 17:11:24 +00:00
parent c2c9a78db9
commit c7a26b7c32
3 changed files with 15 additions and 7 deletions

View File

@ -194,7 +194,7 @@ class E2eKeysHandler(object):
# "unsigned" section # "unsigned" section
for user_id, device_keys in results.items(): for user_id, device_keys in results.items():
for device_id, device_info in device_keys.items(): for device_id, device_info in device_keys.items():
r = json.loads(device_info["key_json"]) r = dict(device_info["keys"])
r["unsigned"] = {} r["unsigned"] = {}
display_name = device_info["device_display_name"] display_name = device_info["device_display_name"]
if display_name is not None: if display_name is not None:

View File

@ -15,6 +15,7 @@
from twisted.internet import defer from twisted.internet import defer
from canonicaljson import encode_canonical_json from canonicaljson import encode_canonical_json
import ujson as json
from ._base import SQLBaseStore from ._base import SQLBaseStore
@ -59,6 +60,7 @@ class EndToEndKeyStore(SQLBaseStore):
"set_e2e_device_keys", _set_e2e_device_keys_txn "set_e2e_device_keys", _set_e2e_device_keys_txn
) )
@defer.inlineCallbacks
def get_e2e_device_keys(self, query_list, include_all_devices=False): def get_e2e_device_keys(self, query_list, include_all_devices=False):
"""Fetch a list of device keys. """Fetch a list of device keys.
Args: Args:
@ -70,13 +72,19 @@ class EndToEndKeyStore(SQLBaseStore):
dict containing "key_json", "device_display_name". dict containing "key_json", "device_display_name".
""" """
if not query_list: if not query_list:
return {} defer.returnValue({})
return self.runInteraction( results = yield self.runInteraction(
"get_e2e_device_keys", self._get_e2e_device_keys_txn, "get_e2e_device_keys", self._get_e2e_device_keys_txn,
query_list, include_all_devices, 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): def _get_e2e_device_keys_txn(self, txn, query_list, include_all_devices):
query_clauses = [] query_clauses = []
query_params = [] query_params = []

View File

@ -33,7 +33,7 @@ class EndToEndKeyStoreTestCase(tests.unittest.TestCase):
@defer.inlineCallbacks @defer.inlineCallbacks
def test_key_without_device_name(self): def test_key_without_device_name(self):
now = 1470174257070 now = 1470174257070
json = '{ "key": "value" }' json = {"key": "value"}
yield self.store.store_device( yield self.store.store_device(
"user", "device", None "user", "device", None
@ -47,14 +47,14 @@ class EndToEndKeyStoreTestCase(tests.unittest.TestCase):
self.assertIn("device", res["user"]) self.assertIn("device", res["user"])
dev = res["user"]["device"] dev = res["user"]["device"]
self.assertDictContainsSubset({ self.assertDictContainsSubset({
"key_json": json, "keys": json,
"device_display_name": None, "device_display_name": None,
}, dev) }, dev)
@defer.inlineCallbacks @defer.inlineCallbacks
def test_get_key_with_device_name(self): def test_get_key_with_device_name(self):
now = 1470174257070 now = 1470174257070
json = '{ "key": "value" }' json = {"key": "value"}
yield self.store.set_e2e_device_keys( yield self.store.set_e2e_device_keys(
"user", "device", now, json) "user", "device", now, json)
@ -67,7 +67,7 @@ class EndToEndKeyStoreTestCase(tests.unittest.TestCase):
self.assertIn("device", res["user"]) self.assertIn("device", res["user"])
dev = res["user"]["device"] dev = res["user"]["device"]
self.assertDictContainsSubset({ self.assertDictContainsSubset({
"key_json": json, "keys": json,
"device_display_name": "display_name", "device_display_name": "display_name",
}, dev) }, dev)