test__base.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154
  1. # Copyright 2015, 2016 OpenMarket Ltd
  2. # Copyright 2019 New Vector 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. import secrets
  16. from typing import Generator, Tuple
  17. from twisted.test.proto_helpers import MemoryReactor
  18. from synapse.server import HomeServer
  19. from synapse.util import Clock
  20. from tests import unittest
  21. class UpdateUpsertManyTests(unittest.HomeserverTestCase):
  22. def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
  23. self.storage = hs.get_datastores().main
  24. self.table_name = "table_" + secrets.token_hex(6)
  25. self.get_success(
  26. self.storage.db_pool.runInteraction(
  27. "create",
  28. lambda x, *a: x.execute(*a),
  29. "CREATE TABLE %s (id INTEGER, username TEXT, value TEXT)"
  30. % (self.table_name,),
  31. )
  32. )
  33. self.get_success(
  34. self.storage.db_pool.runInteraction(
  35. "index",
  36. lambda x, *a: x.execute(*a),
  37. "CREATE UNIQUE INDEX %sindex ON %s(id, username)"
  38. % (self.table_name, self.table_name),
  39. )
  40. )
  41. def _dump_table_to_tuple(self) -> Generator[Tuple[int, str, str], None, None]:
  42. res = self.get_success(
  43. self.storage.db_pool.simple_select_list(
  44. self.table_name, None, ["id, username, value"]
  45. )
  46. )
  47. for i in res:
  48. yield (i["id"], i["username"], i["value"])
  49. def test_upsert_many(self) -> None:
  50. """
  51. Upsert_many will perform the upsert operation across a batch of data.
  52. """
  53. # Add some data to an empty table
  54. key_names = ["id", "username"]
  55. value_names = ["value"]
  56. key_values = [[1, "user1"], [2, "user2"]]
  57. value_values = [["hello"], ["there"]]
  58. self.get_success(
  59. self.storage.db_pool.runInteraction(
  60. "test",
  61. self.storage.db_pool.simple_upsert_many_txn,
  62. self.table_name,
  63. key_names,
  64. key_values,
  65. value_names,
  66. value_values,
  67. )
  68. )
  69. # Check results are what we expect
  70. self.assertEqual(
  71. set(self._dump_table_to_tuple()),
  72. {(1, "user1", "hello"), (2, "user2", "there")},
  73. )
  74. # Update only user2
  75. key_values = [[2, "user2"]]
  76. value_values = [["bleb"]]
  77. self.get_success(
  78. self.storage.db_pool.runInteraction(
  79. "test",
  80. self.storage.db_pool.simple_upsert_many_txn,
  81. self.table_name,
  82. key_names,
  83. key_values,
  84. value_names,
  85. value_values,
  86. )
  87. )
  88. # Check results are what we expect
  89. self.assertEqual(
  90. set(self._dump_table_to_tuple()),
  91. {(1, "user1", "hello"), (2, "user2", "bleb")},
  92. )
  93. def test_simple_update_many(self) -> None:
  94. """
  95. simple_update_many performs many updates at once.
  96. """
  97. # First add some data.
  98. self.get_success(
  99. self.storage.db_pool.simple_insert_many(
  100. table=self.table_name,
  101. keys=("id", "username", "value"),
  102. values=[(1, "alice", "A"), (2, "bob", "B"), (3, "charlie", "C")],
  103. desc="insert",
  104. )
  105. )
  106. # Check the data made it to the table
  107. self.assertEqual(
  108. set(self._dump_table_to_tuple()),
  109. {(1, "alice", "A"), (2, "bob", "B"), (3, "charlie", "C")},
  110. )
  111. # Now use simple_update_many
  112. self.get_success(
  113. self.storage.db_pool.simple_update_many(
  114. table=self.table_name,
  115. key_names=("username",),
  116. key_values=(
  117. ("alice",),
  118. ("bob",),
  119. ("stranger",),
  120. ),
  121. value_names=("value",),
  122. value_values=(
  123. ("aaa!",),
  124. ("bbb!",),
  125. ("???",),
  126. ),
  127. desc="update_many1",
  128. )
  129. )
  130. # Check the table is how we expect:
  131. # charlie has been left alone
  132. self.assertEqual(
  133. set(self._dump_table_to_tuple()),
  134. {(1, "alice", "aaa!"), (2, "bob", "bbb!"), (3, "charlie", "C")},
  135. )