diff --git a/tests/handlers/test_directory.py b/tests/handlers/test_directory.py
index 90aec484c4..367d94eca3 100644
--- a/tests/handlers/test_directory.py
+++ b/tests/handlers/test_directory.py
@@ -13,7 +13,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Any, Awaitable, Callable, Dict
-from unittest.mock import Mock
+from unittest.mock import AsyncMock, Mock
from twisted.test.proto_helpers import MemoryReactor
@@ -27,14 +27,13 @@ from synapse.types import JsonDict, RoomAlias, create_requester
from synapse.util import Clock
from tests import unittest
-from tests.test_utils import make_awaitable
class DirectoryTestCase(unittest.HomeserverTestCase):
"""Tests the directory service."""
def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer:
- self.mock_federation = Mock()
+ self.mock_federation = AsyncMock()
self.mock_registry = Mock()
self.query_handlers: Dict[str, Callable[[dict], Awaitable[JsonDict]]] = {}
@@ -73,9 +72,10 @@ class DirectoryTestCase(unittest.HomeserverTestCase):
self.assertEqual({"room_id": "!8765qwer:test", "servers": ["test"]}, result)
def test_get_remote_association(self) -> None:
- self.mock_federation.make_query.return_value = make_awaitable(
- {"room_id": "!8765qwer:test", "servers": ["test", "remote"]}
- )
+ self.mock_federation.make_query.return_value = {
+ "room_id": "!8765qwer:test",
+ "servers": ["test", "remote"],
+ }
result = self.get_success(self.handler.get_association(self.remote_room))
|
