diff --git a/tests/test_federation.py b/tests/test_federation.py
index 6d15ac7597..779f70467b 100644
--- a/tests/test_federation.py
+++ b/tests/test_federation.py
@@ -13,7 +13,7 @@
# limitations under the License.
from typing import Collection, List, Optional, Union
-from unittest.mock import Mock
+from unittest.mock import AsyncMock, Mock
from twisted.test.proto_helpers import MemoryReactor
@@ -31,7 +31,6 @@ from synapse.util import Clock
from synapse.util.retryutils import NotRetryingDestination
from tests import unittest
-from tests.test_utils import make_awaitable
class MessageAcceptTests(unittest.HomeserverTestCase):
@@ -196,7 +195,7 @@ class MessageAcceptTests(unittest.HomeserverTestCase):
# Register a mock on the store so that the incoming update doesn't fail because
# we don't share a room with the user.
store = self.hs.get_datastores().main
- store.get_rooms_for_user = Mock(return_value=make_awaitable(["!someroom:test"]))
+ store.get_rooms_for_user = AsyncMock(return_value=["!someroom:test"])
# Manually inject a fake device list update. We need this update to include at
# least one prev_id so that the user's device list will need to be retried.
@@ -241,27 +240,24 @@ class MessageAcceptTests(unittest.HomeserverTestCase):
# Register mock device list retrieval on the federation client.
federation_client = self.hs.get_federation_client()
- federation_client.query_user_devices = Mock( # type: ignore[assignment]
- return_value=make_awaitable(
- {
+ federation_client.query_user_devices = AsyncMock( # type: ignore[assignment]
+ return_value={
+ "user_id": remote_user_id,
+ "stream_id": 1,
+ "devices": [],
+ "master_key": {
"user_id": remote_user_id,
- "stream_id": 1,
- "devices": [],
- "master_key": {
- "user_id": remote_user_id,
- "usage": ["master"],
- "keys": {"ed25519:" + remote_master_key: remote_master_key},
- },
- "self_signing_key": {
- "user_id": remote_user_id,
- "usage": ["self_signing"],
- "keys": {
- "ed25519:"
- + remote_self_signing_key: remote_self_signing_key
- },
+ "usage": ["master"],
+ "keys": {"ed25519:" + remote_master_key: remote_master_key},
+ },
+ "self_signing_key": {
+ "user_id": remote_user_id,
+ "usage": ["self_signing"],
+ "keys": {
+ "ed25519:" + remote_self_signing_key: remote_self_signing_key
},
- }
- )
+ },
+ }
)
# Resync the device list.
|
