1
0

test_base.py 7.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211
  1. # Copyright 2014-2016 OpenMarket Ltd
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from collections import OrderedDict
  15. from typing import Generator
  16. from unittest.mock import Mock
  17. from twisted.internet import defer
  18. from synapse.storage._base import SQLBaseStore
  19. from synapse.storage.database import DatabasePool
  20. from synapse.storage.engines import create_engine
  21. from tests import unittest
  22. from tests.server import TestHomeServer
  23. from tests.utils import default_config
  24. class SQLBaseStoreTestCase(unittest.TestCase):
  25. """Test the "simple" SQL generating methods in SQLBaseStore."""
  26. def setUp(self) -> None:
  27. self.db_pool = Mock(spec=["runInteraction"])
  28. self.mock_txn = Mock()
  29. self.mock_conn = Mock(spec_set=["cursor", "rollback", "commit"])
  30. self.mock_conn.cursor.return_value = self.mock_txn
  31. self.mock_conn.rollback.return_value = None
  32. # Our fake runInteraction just runs synchronously inline
  33. def runInteraction(func, *args, **kwargs) -> defer.Deferred: # type: ignore[no-untyped-def]
  34. return defer.succeed(func(self.mock_txn, *args, **kwargs))
  35. self.db_pool.runInteraction = runInteraction
  36. def runWithConnection(func, *args, **kwargs): # type: ignore[no-untyped-def]
  37. return defer.succeed(func(self.mock_conn, *args, **kwargs))
  38. self.db_pool.runWithConnection = runWithConnection
  39. config = default_config(name="test", parse=True)
  40. hs = TestHomeServer("test", config=config)
  41. sqlite_config = {"name": "sqlite3"}
  42. engine = create_engine(sqlite_config)
  43. fake_engine = Mock(wraps=engine)
  44. fake_engine.in_transaction.return_value = False
  45. db = DatabasePool(Mock(), Mock(config=sqlite_config), fake_engine)
  46. db._db_pool = self.db_pool
  47. self.datastore = SQLBaseStore(db, None, hs) # type: ignore[arg-type]
  48. @defer.inlineCallbacks
  49. def test_insert_1col(self) -> Generator["defer.Deferred[object]", object, None]:
  50. self.mock_txn.rowcount = 1
  51. yield defer.ensureDeferred(
  52. self.datastore.db_pool.simple_insert(
  53. table="tablename", values={"columname": "Value"}
  54. )
  55. )
  56. self.mock_txn.execute.assert_called_with(
  57. "INSERT INTO tablename (columname) VALUES(?)", ("Value",)
  58. )
  59. @defer.inlineCallbacks
  60. def test_insert_3cols(self) -> Generator["defer.Deferred[object]", object, None]:
  61. self.mock_txn.rowcount = 1
  62. yield defer.ensureDeferred(
  63. self.datastore.db_pool.simple_insert(
  64. table="tablename",
  65. # Use OrderedDict() so we can assert on the SQL generated
  66. values=OrderedDict([("colA", 1), ("colB", 2), ("colC", 3)]),
  67. )
  68. )
  69. self.mock_txn.execute.assert_called_with(
  70. "INSERT INTO tablename (colA, colB, colC) VALUES(?, ?, ?)", (1, 2, 3)
  71. )
  72. @defer.inlineCallbacks
  73. def test_select_one_1col(self) -> Generator["defer.Deferred[object]", object, None]:
  74. self.mock_txn.rowcount = 1
  75. self.mock_txn.__iter__ = Mock(return_value=iter([("Value",)]))
  76. value = yield defer.ensureDeferred(
  77. self.datastore.db_pool.simple_select_one_onecol(
  78. table="tablename", keyvalues={"keycol": "TheKey"}, retcol="retcol"
  79. )
  80. )
  81. self.assertEqual("Value", value)
  82. self.mock_txn.execute.assert_called_with(
  83. "SELECT retcol FROM tablename WHERE keycol = ?", ["TheKey"]
  84. )
  85. @defer.inlineCallbacks
  86. def test_select_one_3col(self) -> Generator["defer.Deferred[object]", object, None]:
  87. self.mock_txn.rowcount = 1
  88. self.mock_txn.fetchone.return_value = (1, 2, 3)
  89. ret = yield defer.ensureDeferred(
  90. self.datastore.db_pool.simple_select_one(
  91. table="tablename",
  92. keyvalues={"keycol": "TheKey"},
  93. retcols=["colA", "colB", "colC"],
  94. )
  95. )
  96. self.assertEqual({"colA": 1, "colB": 2, "colC": 3}, ret)
  97. self.mock_txn.execute.assert_called_with(
  98. "SELECT colA, colB, colC FROM tablename WHERE keycol = ?", ["TheKey"]
  99. )
  100. @defer.inlineCallbacks
  101. def test_select_one_missing(
  102. self,
  103. ) -> Generator["defer.Deferred[object]", object, None]:
  104. self.mock_txn.rowcount = 0
  105. self.mock_txn.fetchone.return_value = None
  106. ret = yield defer.ensureDeferred(
  107. self.datastore.db_pool.simple_select_one(
  108. table="tablename",
  109. keyvalues={"keycol": "Not here"},
  110. retcols=["colA"],
  111. allow_none=True,
  112. )
  113. )
  114. self.assertFalse(ret)
  115. @defer.inlineCallbacks
  116. def test_select_list(self) -> Generator["defer.Deferred[object]", object, None]:
  117. self.mock_txn.rowcount = 3
  118. self.mock_txn.__iter__ = Mock(return_value=iter([(1,), (2,), (3,)]))
  119. self.mock_txn.description = (("colA", None, None, None, None, None, None),)
  120. ret = yield defer.ensureDeferred(
  121. self.datastore.db_pool.simple_select_list(
  122. table="tablename", keyvalues={"keycol": "A set"}, retcols=["colA"]
  123. )
  124. )
  125. self.assertEqual([{"colA": 1}, {"colA": 2}, {"colA": 3}], ret)
  126. self.mock_txn.execute.assert_called_with(
  127. "SELECT colA FROM tablename WHERE keycol = ?", ["A set"]
  128. )
  129. @defer.inlineCallbacks
  130. def test_update_one_1col(self) -> Generator["defer.Deferred[object]", object, None]:
  131. self.mock_txn.rowcount = 1
  132. yield defer.ensureDeferred(
  133. self.datastore.db_pool.simple_update_one(
  134. table="tablename",
  135. keyvalues={"keycol": "TheKey"},
  136. updatevalues={"columnname": "New Value"},
  137. )
  138. )
  139. self.mock_txn.execute.assert_called_with(
  140. "UPDATE tablename SET columnname = ? WHERE keycol = ?",
  141. ["New Value", "TheKey"],
  142. )
  143. @defer.inlineCallbacks
  144. def test_update_one_4cols(
  145. self,
  146. ) -> Generator["defer.Deferred[object]", object, None]:
  147. self.mock_txn.rowcount = 1
  148. yield defer.ensureDeferred(
  149. self.datastore.db_pool.simple_update_one(
  150. table="tablename",
  151. keyvalues=OrderedDict([("colA", 1), ("colB", 2)]),
  152. updatevalues=OrderedDict([("colC", 3), ("colD", 4)]),
  153. )
  154. )
  155. self.mock_txn.execute.assert_called_with(
  156. "UPDATE tablename SET colC = ?, colD = ? WHERE" " colA = ? AND colB = ?",
  157. [3, 4, 1, 2],
  158. )
  159. @defer.inlineCallbacks
  160. def test_delete_one(self) -> Generator["defer.Deferred[object]", object, None]:
  161. self.mock_txn.rowcount = 1
  162. yield defer.ensureDeferred(
  163. self.datastore.db_pool.simple_delete_one(
  164. table="tablename", keyvalues={"keycol": "Go away"}
  165. )
  166. )
  167. self.mock_txn.execute.assert_called_with(
  168. "DELETE FROM tablename WHERE keycol = ?", ["Go away"]
  169. )