From 91fa69e0299167dad4df9831b8d175a99564b266 Mon Sep 17 00:00:00 2001
From: Richard van der Hoff <richard@matrix.org>
Date: Wed, 3 Aug 2016 11:48:32 +0100
Subject: [PATCH] keys/query: return all users which were asked for

In the situation where all of a user's devices get deleted, we want to
indicate this to a client, so we want to return an empty dictionary, rather
than nothing at all.
---
 synapse/handlers/e2e_keys.py    |  9 ++++---
 tests/handlers/test_e2e_keys.py | 46 +++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 3 deletions(-)
 create mode 100644 tests/handlers/test_e2e_keys.py

diff --git a/synapse/handlers/e2e_keys.py b/synapse/handlers/e2e_keys.py
index 1312cdf5a..950fc927b 100644
--- a/synapse/handlers/e2e_keys.py
+++ b/synapse/handlers/e2e_keys.py
@@ -99,6 +99,7 @@ class E2eKeysHandler(object):
         """
         local_query = []
 
+        result_dict = {}
         for user_id, device_ids in query.items():
             if not self.is_mine_id(user_id):
                 logger.warning("Request for keys for non-local user %s",
@@ -111,15 +112,17 @@ class E2eKeysHandler(object):
                 for device_id in device_ids:
                     local_query.append((user_id, device_id))
 
+            # make sure that each queried user appears in the result dict
+            result_dict[user_id] = {}
+
         results = yield self.store.get_e2e_device_keys(local_query)
 
         # un-jsonify the results
-        json_result = collections.defaultdict(dict)
         for user_id, device_keys in results.items():
             for device_id, json_bytes in device_keys.items():
-                json_result[user_id][device_id] = json.loads(json_bytes)
+                result_dict[user_id][device_id] = json.loads(json_bytes)
 
-        defer.returnValue(json_result)
+        defer.returnValue(result_dict)
 
     @defer.inlineCallbacks
     def on_federation_query_client_keys(self, query_body):
diff --git a/tests/handlers/test_e2e_keys.py b/tests/handlers/test_e2e_keys.py
new file mode 100644
index 000000000..878a54dc3
--- /dev/null
+++ b/tests/handlers/test_e2e_keys.py
@@ -0,0 +1,46 @@
+# -*- coding: utf-8 -*-
+# Copyright 2016 OpenMarket Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import mock
+from twisted.internet import defer
+
+import synapse.api.errors
+import synapse.handlers.e2e_keys
+
+import synapse.storage
+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.handler = None  # type: synapse.handlers.e2e_keys.E2eKeysHandler
+
+    @defer.inlineCallbacks
+    def setUp(self):
+        self.hs = yield utils.setup_test_homeserver(
+            handlers=None,
+            replication_layer=mock.Mock(),
+        )
+        self.handler = synapse.handlers.e2e_keys.E2eKeysHandler(self.hs)
+
+    @defer.inlineCallbacks
+    def test_query_local_devices_no_devices(self):
+        """If the user has no devices, we expect an empty list.
+        """
+        local_user = "@boris:" + self.hs.hostname
+        res = yield self.handler.query_local_devices({local_user: None})
+        self.assertDictEqual(res, {local_user: {}})