diff --git a/tests/federation/test_complexity.py b/tests/federation/test_complexity.py
index b8ca118716..3d880c499d 100644
--- a/tests/federation/test_complexity.py
+++ b/tests/federation/test_complexity.py
@@ -15,8 +15,6 @@
from mock import Mock
-from twisted.internet import defer
-
from synapse.api.errors import Codes, SynapseError
from synapse.rest import admin
from synapse.rest.client.v1 import login, room
@@ -60,7 +58,7 @@ class RoomComplexityTests(unittest.FederatingHomeserverTestCase):
# Artificially raise the complexity
store = self.hs.get_datastore()
- store.get_current_state_event_counts = lambda x: defer.succeed(500 * 1.23)
+ store.get_current_state_event_counts = lambda x: make_awaitable(500 * 1.23)
# Get the room complexity again -- make sure it's our artificial value
request, channel = self.make_request(
@@ -79,9 +77,11 @@ class RoomComplexityTests(unittest.FederatingHomeserverTestCase):
fed_transport = self.hs.get_federation_transport_client()
# Mock out some things, because we don't want to test the whole join
- fed_transport.client.get_json = Mock(return_value=make_awaitable({"v1": 9999}))
+ fed_transport.client.get_json = Mock(
+ side_effect=lambda *args, **kwargs: make_awaitable({"v1": 9999})
+ )
handler.federation_handler.do_invite_join = Mock(
- return_value=make_awaitable(("", 1))
+ side_effect=lambda *args, **kwargs: make_awaitable(("", 1))
)
d = handler._remote_join(
@@ -110,9 +110,11 @@ class RoomComplexityTests(unittest.FederatingHomeserverTestCase):
fed_transport = self.hs.get_federation_transport_client()
# Mock out some things, because we don't want to test the whole join
- fed_transport.client.get_json = Mock(return_value=make_awaitable({"v1": 9999}))
+ fed_transport.client.get_json = Mock(
+ side_effect=lambda *args, **kwargs: make_awaitable({"v1": 9999})
+ )
handler.federation_handler.do_invite_join = Mock(
- return_value=make_awaitable(("", 1))
+ side_effect=lambda *args, **kwargs: make_awaitable(("", 1))
)
d = handler._remote_join(
@@ -148,13 +150,15 @@ class RoomComplexityTests(unittest.FederatingHomeserverTestCase):
fed_transport = self.hs.get_federation_transport_client()
# Mock out some things, because we don't want to test the whole join
- fed_transport.client.get_json = Mock(return_value=make_awaitable(None))
+ fed_transport.client.get_json = Mock(
+ side_effect=lambda *args, **kwargs: make_awaitable(None)
+ )
handler.federation_handler.do_invite_join = Mock(
- return_value=make_awaitable(("", 1))
+ side_effect=lambda *args, **kwargs: make_awaitable(("", 1))
)
# Artificially raise the complexity
- self.hs.get_datastore().get_current_state_event_counts = lambda x: defer.succeed(
+ self.hs.get_datastore().get_current_state_event_counts = lambda x: make_awaitable(
600
)
@@ -204,9 +208,11 @@ class RoomComplexityAdminTests(unittest.FederatingHomeserverTestCase):
fed_transport = self.hs.get_federation_transport_client()
# Mock out some things, because we don't want to test the whole join
- fed_transport.client.get_json = Mock(return_value=make_awaitable({"v1": 9999}))
+ fed_transport.client.get_json = Mock(
+ side_effect=lambda *args, **kwargs: make_awaitable({"v1": 9999})
+ )
handler.federation_handler.do_invite_join = Mock(
- return_value=make_awaitable(("", 1))
+ side_effect=lambda *args, **kwargs: make_awaitable(("", 1))
)
d = handler._remote_join(
@@ -234,9 +240,11 @@ class RoomComplexityAdminTests(unittest.FederatingHomeserverTestCase):
fed_transport = self.hs.get_federation_transport_client()
# Mock out some things, because we don't want to test the whole join
- fed_transport.client.get_json = Mock(return_value=make_awaitable({"v1": 9999}))
+ fed_transport.client.get_json = Mock(
+ side_effect=lambda *args, **kwargs: make_awaitable({"v1": 9999})
+ )
handler.federation_handler.do_invite_join = Mock(
- return_value=make_awaitable(("", 1))
+ side_effect=lambda *args, **kwargs: make_awaitable(("", 1))
)
d = handler._remote_join(
|
