make count_monthly_users async synapse/handlers/auth.py

2 files changed, 28 insertions, 21 deletions

diff --git a/tests/handlers/test_auth.py b/tests/handlers/test_auth.py
index e01f14a10a..440a453082 100644
--- a/tests/handlers/test_auth.py
+++ b/tests/handlers/test_auth.py
@@ -77,38 +77,37 @@ class AuthTestCase(unittest.TestCase):
         v.satisfy_general(verify_nonce)
         v.verify(macaroon, self.hs.config.macaroon_secret_key)
 
+    @defer.inlineCallbacks
     def test_short_term_login_token_gives_user_id(self):
         self.hs.clock.now = 1000
 
         token = self.macaroon_generator.generate_short_term_login_token(
             "a_user", 5000
         )
-
-        self.assertEqual(
-            "a_user",
-            self.auth_handler.validate_short_term_login_token_and_get_user_id(
-                token
-            )
+        user_id = yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
+            token
         )
+        self.assertEqual("a_user", user_id)
 
         # when we advance the clock, the token should be rejected
         self.hs.clock.now = 6000
         with self.assertRaises(synapse.api.errors.AuthError):
-            self.auth_handler.validate_short_term_login_token_and_get_user_id(
+            yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
                 token
             )
 
+    @defer.inlineCallbacks
     def test_short_term_login_token_cannot_replace_user_id(self):
         token = self.macaroon_generator.generate_short_term_login_token(
             "a_user", 5000
         )
         macaroon = pymacaroons.Macaroon.deserialize(token)
 
+        user_id = yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
+            macaroon.serialize()
+        )
         self.assertEqual(
-            "a_user",
-            self.auth_handler.validate_short_term_login_token_and_get_user_id(
-                macaroon.serialize()
-            )
+            "a_user", user_id
         )
 
         # add another "user_id" caveat, which might allow us to override the
@@ -116,7 +115,7 @@ class AuthTestCase(unittest.TestCase):
         macaroon.add_first_party_caveat("user_id = b_user")
 
         with self.assertRaises(synapse.api.errors.AuthError):
-            self.auth_handler.validate_short_term_login_token_and_get_user_id(
+            yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
                 macaroon.serialize()
             )
 
@@ -126,7 +125,7 @@ class AuthTestCase(unittest.TestCase):
         # Ensure does not throw exception
         yield self.auth_handler.get_access_token_for_user_id('user_a')
 
-        self.auth_handler.validate_short_term_login_token_and_get_user_id(
+        yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
             self._get_macaroon().serialize()
         )
 
@@ -134,24 +133,30 @@ class AuthTestCase(unittest.TestCase):
     def test_mau_limits_exceeded(self):
         self.hs.config.limit_usage_by_mau = True
         self.hs.get_datastore().count_monthly_users = Mock(
-            return_value=self.large_number_of_users
+            return_value=defer.succeed(self.large_number_of_users)
         )
+
         with self.assertRaises(AuthError):
             yield self.auth_handler.get_access_token_for_user_id('user_a')
+
+        self.hs.get_datastore().count_monthly_users = Mock(
+            return_value=defer.succeed(self.large_number_of_users)
+        )
         with self.assertRaises(AuthError):
-            self.auth_handler.validate_short_term_login_token_and_get_user_id(
+            yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
                 self._get_macaroon().serialize()
             )
 
     @defer.inlineCallbacks
     def test_mau_limits_not_exceeded(self):
         self.hs.config.limit_usage_by_mau = True
+
         self.hs.get_datastore().count_monthly_users = Mock(
-            return_value=self.small_number_of_users
+            return_value=defer.succeed(self.small_number_of_users)
         )
         # Ensure does not raise exception
         yield self.auth_handler.get_access_token_for_user_id('user_a')
-        self.auth_handler.validate_short_term_login_token_and_get_user_id(
+        yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
             self._get_macaroon().serialize()
         )
 
diff --git a/tests/handlers/test_register.py b/tests/handlers/test_register.py
index a5a8e7c954..0937d71cf6 100644
--- a/tests/handlers/test_register.py
+++ b/tests/handlers/test_register.py
@@ -90,7 +90,7 @@ class RegistrationTestCase(unittest.TestCase):
         lots_of_users = 100
         small_number_users = 1
 
-        store.count_monthly_users = Mock(return_value=lots_of_users)
+        store.count_monthly_users = Mock(return_value=defer.succeed(lots_of_users))
 
         # Ensure does not throw exception
         yield self.handler.get_or_create_user(requester, 'a', display_name)
@@ -100,7 +100,7 @@ class RegistrationTestCase(unittest.TestCase):
         with self.assertRaises(RegistrationError):
             yield self.handler.get_or_create_user(requester, 'b', display_name)
 
-        store.count_monthly_users = Mock(return_value=small_number_users)
+        store.count_monthly_users = Mock(return_value=defer.succeed(small_number_users))
 
         self._macaroon_mock_generator("another_secret")
 
@@ -108,12 +108,14 @@ class RegistrationTestCase(unittest.TestCase):
         yield self.handler.get_or_create_user("@neil:matrix.org", 'c', "Neil")
 
         self._macaroon_mock_generator("another another secret")
-        store.count_monthly_users = Mock(return_value=lots_of_users)
+        store.count_monthly_users = Mock(return_value=defer.succeed(lots_of_users))
+
         with self.assertRaises(RegistrationError):
             yield self.handler.register(localpart=local_part)
 
         self._macaroon_mock_generator("another another secret")
-        store.count_monthly_users = Mock(return_value=lots_of_users)
+        store.count_monthly_users = Mock(return_value=defer.succeed(lots_of_users))
+
         with self.assertRaises(RegistrationError):
             yield self.handler.register_saml2(local_part)