diff options
Diffstat (limited to 'tests/handlers/test_auth.py')
| -rw-r--r-- | tests/handlers/test_auth.py | 21 |
1 files changed, 11 insertions, 10 deletions
diff --git a/tests/handlers/test_auth.py b/tests/handlers/test_auth.py index c01b04e1dc..4b3fb018b1 100644 --- a/tests/handlers/test_auth.py +++ b/tests/handlers/test_auth.py @@ -24,6 +24,7 @@ from synapse.api.errors import ResourceLimitError from synapse.handlers.auth import AuthHandler from tests import unittest +from tests.test_utils import make_awaitable from tests.utils import setup_test_homeserver @@ -142,7 +143,7 @@ class AuthTestCase(unittest.TestCase): def test_mau_limits_exceeded_large(self): self.auth_blocking._limit_usage_by_mau = True self.hs.get_datastore().get_monthly_active_count = Mock( - return_value=defer.succeed(self.large_number_of_users) + side_effect=lambda: make_awaitable(self.large_number_of_users) ) with self.assertRaises(ResourceLimitError): @@ -153,7 +154,7 @@ class AuthTestCase(unittest.TestCase): ) self.hs.get_datastore().get_monthly_active_count = Mock( - return_value=defer.succeed(self.large_number_of_users) + side_effect=lambda: make_awaitable(self.large_number_of_users) ) with self.assertRaises(ResourceLimitError): yield defer.ensureDeferred( @@ -168,7 +169,7 @@ class AuthTestCase(unittest.TestCase): # If not in monthly active cohort self.hs.get_datastore().get_monthly_active_count = Mock( - return_value=defer.succeed(self.auth_blocking._max_mau_value) + side_effect=lambda: make_awaitable(self.auth_blocking._max_mau_value) ) with self.assertRaises(ResourceLimitError): yield defer.ensureDeferred( @@ -178,7 +179,7 @@ class AuthTestCase(unittest.TestCase): ) self.hs.get_datastore().get_monthly_active_count = Mock( - return_value=defer.succeed(self.auth_blocking._max_mau_value) + side_effect=lambda: make_awaitable(self.auth_blocking._max_mau_value) ) with self.assertRaises(ResourceLimitError): yield defer.ensureDeferred( @@ -188,10 +189,10 @@ class AuthTestCase(unittest.TestCase): ) # If in monthly active cohort self.hs.get_datastore().user_last_seen_monthly_active = Mock( - return_value=defer.succeed(self.hs.get_clock().time_msec()) + side_effect=lambda user_id: make_awaitable(self.hs.get_clock().time_msec()) ) self.hs.get_datastore().get_monthly_active_count = Mock( - return_value=defer.succeed(self.auth_blocking._max_mau_value) + side_effect=lambda: make_awaitable(self.auth_blocking._max_mau_value) ) yield defer.ensureDeferred( self.auth_handler.get_access_token_for_user_id( @@ -199,10 +200,10 @@ class AuthTestCase(unittest.TestCase): ) ) self.hs.get_datastore().user_last_seen_monthly_active = Mock( - return_value=defer.succeed(self.hs.get_clock().time_msec()) + side_effect=lambda user_id: make_awaitable(self.hs.get_clock().time_msec()) ) self.hs.get_datastore().get_monthly_active_count = Mock( - return_value=defer.succeed(self.auth_blocking._max_mau_value) + side_effect=lambda: make_awaitable(self.auth_blocking._max_mau_value) ) yield defer.ensureDeferred( self.auth_handler.validate_short_term_login_token_and_get_user_id( @@ -215,7 +216,7 @@ class AuthTestCase(unittest.TestCase): self.auth_blocking._limit_usage_by_mau = True self.hs.get_datastore().get_monthly_active_count = Mock( - return_value=defer.succeed(self.small_number_of_users) + side_effect=lambda: make_awaitable(self.small_number_of_users) ) # Ensure does not raise exception yield defer.ensureDeferred( @@ -225,7 +226,7 @@ class AuthTestCase(unittest.TestCase): ) self.hs.get_datastore().get_monthly_active_count = Mock( - return_value=defer.succeed(self.small_number_of_users) + side_effect=lambda: make_awaitable(self.small_number_of_users) ) yield defer.ensureDeferred( self.auth_handler.validate_short_term_login_token_and_get_user_id( |