test_end_to_end_keys.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596
  1. # -*- coding: utf-8 -*-
  2. # Copyright 2016 OpenMarket Ltd
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. # You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. from twisted.internet import defer
  16. import tests.unittest
  17. import tests.utils
  18. class EndToEndKeyStoreTestCase(tests.unittest.TestCase):
  19. @defer.inlineCallbacks
  20. def setUp(self):
  21. hs = yield tests.utils.setup_test_homeserver(self.addCleanup)
  22. self.store = hs.get_datastore()
  23. @defer.inlineCallbacks
  24. def test_key_without_device_name(self):
  25. now = 1470174257070
  26. json = {"key": "value"}
  27. yield self.store.store_device("user", "device", None)
  28. yield self.store.set_e2e_device_keys("user", "device", now, json)
  29. res = yield self.store.get_e2e_device_keys((("user", "device"),))
  30. self.assertIn("user", res)
  31. self.assertIn("device", res["user"])
  32. dev = res["user"]["device"]
  33. self.assertDictContainsSubset(json, dev)
  34. @defer.inlineCallbacks
  35. def test_reupload_key(self):
  36. now = 1470174257070
  37. json = {"key": "value"}
  38. yield self.store.store_device("user", "device", None)
  39. changed = yield self.store.set_e2e_device_keys("user", "device", now, json)
  40. self.assertTrue(changed)
  41. # If we try to upload the same key then we should be told nothing
  42. # changed
  43. changed = yield self.store.set_e2e_device_keys("user", "device", now, json)
  44. self.assertFalse(changed)
  45. @defer.inlineCallbacks
  46. def test_get_key_with_device_name(self):
  47. now = 1470174257070
  48. json = {"key": "value"}
  49. yield self.store.set_e2e_device_keys("user", "device", now, json)
  50. yield self.store.store_device("user", "device", "display_name")
  51. res = yield self.store.get_e2e_device_keys((("user", "device"),))
  52. self.assertIn("user", res)
  53. self.assertIn("device", res["user"])
  54. dev = res["user"]["device"]
  55. self.assertDictContainsSubset(
  56. {"key": "value", "unsigned": {"device_display_name": "display_name"}}, dev
  57. )
  58. @defer.inlineCallbacks
  59. def test_multiple_devices(self):
  60. now = 1470174257070
  61. yield self.store.store_device("user1", "device1", None)
  62. yield self.store.store_device("user1", "device2", None)
  63. yield self.store.store_device("user2", "device1", None)
  64. yield self.store.store_device("user2", "device2", None)
  65. yield self.store.set_e2e_device_keys("user1", "device1", now, {"key": "json11"})
  66. yield self.store.set_e2e_device_keys("user1", "device2", now, {"key": "json12"})
  67. yield self.store.set_e2e_device_keys("user2", "device1", now, {"key": "json21"})
  68. yield self.store.set_e2e_device_keys("user2", "device2", now, {"key": "json22"})
  69. res = yield self.store.get_e2e_device_keys(
  70. (("user1", "device1"), ("user2", "device2"))
  71. )
  72. self.assertIn("user1", res)
  73. self.assertIn("device1", res["user1"])
  74. self.assertNotIn("device2", res["user1"])
  75. self.assertIn("user2", res)
  76. self.assertNotIn("device1", res["user2"])
  77. self.assertIn("device2", res["user2"])