test__base.py 4.9 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, List, Tuple, cast
  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. yield from cast(
  43. List[Tuple[int, str, str]],
  44. self.get_success(
  45. self.storage.db_pool.simple_select_list(
  46. self.table_name, None, ["id, username, value"]
  47. )
  48. ),
  49. )
  50. def test_upsert_many(self) -> None:
  51. """
  52. Upsert_many will perform the upsert operation across a batch of data.
  53. """
  54. # Add some data to an empty table
  55. key_names = ["id", "username"]
  56. value_names = ["value"]
  57. key_values = [[1, "user1"], [2, "user2"]]
  58. value_values = [["hello"], ["there"]]
  59. self.get_success(
  60. self.storage.db_pool.runInteraction(
  61. "test",
  62. self.storage.db_pool.simple_upsert_many_txn,
  63. self.table_name,
  64. key_names,
  65. key_values,
  66. value_names,
  67. value_values,
  68. )
  69. )
  70. # Check results are what we expect
  71. self.assertEqual(
  72. set(self._dump_table_to_tuple()),
  73. {(1, "user1", "hello"), (2, "user2", "there")},
  74. )
  75. # Update only user2
  76. key_values = [[2, "user2"]]
  77. value_values = [["bleb"]]
  78. self.get_success(
  79. self.storage.db_pool.runInteraction(
  80. "test",
  81. self.storage.db_pool.simple_upsert_many_txn,
  82. self.table_name,
  83. key_names,
  84. key_values,
  85. value_names,
  86. value_values,
  87. )
  88. )
  89. # Check results are what we expect
  90. self.assertEqual(
  91. set(self._dump_table_to_tuple()),
  92. {(1, "user1", "hello"), (2, "user2", "bleb")},
  93. )
  94. def test_simple_update_many(self) -> None:
  95. """
  96. simple_update_many performs many updates at once.
  97. """
  98. # First add some data.
  99. self.get_success(
  100. self.storage.db_pool.simple_insert_many(
  101. table=self.table_name,
  102. keys=("id", "username", "value"),
  103. values=[(1, "alice", "A"), (2, "bob", "B"), (3, "charlie", "C")],
  104. desc="insert",
  105. )
  106. )
  107. # Check the data made it to the table
  108. self.assertEqual(
  109. set(self._dump_table_to_tuple()),
  110. {(1, "alice", "A"), (2, "bob", "B"), (3, "charlie", "C")},
  111. )
  112. # Now use simple_update_many
  113. self.get_success(
  114. self.storage.db_pool.simple_update_many(
  115. table=self.table_name,
  116. key_names=("username",),
  117. key_values=(
  118. ("alice",),
  119. ("bob",),
  120. ("stranger",),
  121. ),
  122. value_names=("value",),
  123. value_values=(
  124. ("aaa!",),
  125. ("bbb!",),
  126. ("???",),
  127. ),
  128. desc="update_many1",
  129. )
  130. )
  131. # Check the table is how we expect:
  132. # charlie has been left alone
  133. self.assertEqual(
  134. set(self._dump_table_to_tuple()),
  135. {(1, "alice", "aaa!"), (2, "bob", "bbb!"), (3, "charlie", "C")},
  136. )