diff --git a/tests/storage/test_account_data.py b/tests/storage/test_account_data.py
index 01af49a16b..d697d2bc1e 100644
--- a/tests/storage/test_account_data.py
+++ b/tests/storage/test_account_data.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Iterable, Set
+from typing import Iterable, Optional, Set
from synapse.api.constants import AccountDataTypes
@@ -25,7 +25,7 @@ class IgnoredUsersTestCase(unittest.HomeserverTestCase):
self.user = "@user:test"
def _update_ignore_list(
- self, *ignored_user_ids: Iterable[str], ignorer_user_id: str = None
+ self, *ignored_user_ids: Iterable[str], ignorer_user_id: Optional[str] = None
) -> None:
"""Update the account data to block the given users."""
if ignorer_user_id is None:
diff --git a/tests/storage/test_background_update.py b/tests/storage/test_background_update.py
index d77c001506..6156dfac4e 100644
--- a/tests/storage/test_background_update.py
+++ b/tests/storage/test_background_update.py
@@ -12,15 +12,14 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-# Use backported mock for AsyncMock support on Python 3.6.
-from mock import Mock
+from unittest.mock import Mock
from twisted.internet.defer import Deferred, ensureDeferred
from synapse.storage.background_updates import BackgroundUpdater
from tests import unittest
-from tests.test_utils import make_awaitable
+from tests.test_utils import make_awaitable, simple_async_mock
class BackgroundUpdateTestCase(unittest.HomeserverTestCase):
@@ -116,14 +115,14 @@ class BackgroundUpdateControllerTestCase(unittest.HomeserverTestCase):
)
# Mock out the AsyncContextManager
- self._update_ctx_manager = Mock(spec=["__aenter__", "__aexit__"])
- self._update_ctx_manager.__aenter__ = Mock(
- return_value=make_awaitable(None),
- )
- self._update_ctx_manager.__aexit__ = Mock(return_value=make_awaitable(None))
+ class MockCM:
+ __aenter__ = simple_async_mock(return_value=None)
+ __aexit__ = simple_async_mock(return_value=None)
+
+ self._update_ctx_manager = MockCM
# Mock out the `update_handler` callback
- self._on_update = Mock(return_value=self._update_ctx_manager)
+ self._on_update = Mock(return_value=self._update_ctx_manager())
# Define a default batch size value that's not the same as the internal default
# value (100).
diff --git a/tests/storage/test_base.py b/tests/storage/test_base.py
index ddad44bd6c..3e4f0579c9 100644
--- a/tests/storage/test_base.py
+++ b/tests/storage/test_base.py
@@ -23,7 +23,8 @@ from synapse.storage.database import DatabasePool
from synapse.storage.engines import create_engine
from tests import unittest
-from tests.utils import TestHomeServer, default_config
+from tests.server import TestHomeServer
+from tests.utils import default_config
class SQLBaseStoreTestCase(unittest.TestCase):
diff --git a/tests/storage/test_e2e_room_keys.py b/tests/storage/test_e2e_room_keys.py
index 9b6b425425..7556171d8a 100644
--- a/tests/storage/test_e2e_room_keys.py
+++ b/tests/storage/test_e2e_room_keys.py
@@ -12,10 +12,12 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from synapse.storage.databases.main.e2e_room_keys import RoomKey
+
from tests import unittest
# sample room_key data for use in the tests
-room_key = {
+room_key: RoomKey = {
"first_message_index": 1,
"forwarded_count": 1,
"is_verified": False,
diff --git a/tests/storage/test_event_federation.py b/tests/storage/test_event_federation.py
index c3fcf7e7b4..ecfda7677e 100644
--- a/tests/storage/test_event_federation.py
+++ b/tests/storage/test_event_federation.py
@@ -550,7 +550,7 @@ class EventFederationWorkerStoreTestCase(tests.unittest.HomeserverTestCase):
self.store.db_pool.simple_select_one_onecol(
table="federation_inbound_events_staging",
keyvalues={"room_id": room_id},
- retcol="COALESCE(COUNT(*), 0)",
+ retcol="COUNT(*)",
desc="test_prune_inbound_federation_queue",
)
)
diff --git a/tests/storage/test_event_push_actions.py b/tests/storage/test_event_push_actions.py
index bb5939ba4a..738f3ad1dc 100644
--- a/tests/storage/test_event_push_actions.py
+++ b/tests/storage/test_event_push_actions.py
@@ -14,6 +14,8 @@
from unittest.mock import Mock
+from synapse.storage.databases.main.event_push_actions import NotifCounts
+
from tests.unittest import HomeserverTestCase
USER_ID = "@user:example.com"
@@ -57,11 +59,11 @@ class EventPushActionsStoreTestCase(HomeserverTestCase):
)
self.assertEquals(
counts,
- {
- "notify_count": noitf_count,
- "unread_count": 0, # Unread counts are tested in the sync tests.
- "highlight_count": highlight_count,
- },
+ NotifCounts(
+ notify_count=noitf_count,
+ unread_count=0, # Unread counts are tested in the sync tests.
+ highlight_count=highlight_count,
+ ),
)
def _inject_actions(stream, action):
diff --git a/tests/storage/test_roommember.py b/tests/storage/test_roommember.py
index fccab733c0..5cfdfe9b85 100644
--- a/tests/storage/test_roommember.py
+++ b/tests/storage/test_roommember.py
@@ -19,8 +19,8 @@ from synapse.rest.client import login, room
from synapse.types import UserID, create_requester
from tests import unittest
+from tests.server import TestHomeServer
from tests.test_utils import event_injection
-from tests.utils import TestHomeServer
class RoomMemberStoreTestCase(unittest.HomeserverTestCase):
|
