Convert auth handler to async/await (#7261)

2 files changed, 55 insertions, 29 deletions

diff --git a/tests/handlers/test_auth.py b/tests/handlers/test_auth.py
index b03103d96f..52c4ac8b11 100644
--- a/tests/handlers/test_auth.py
+++ b/tests/handlers/test_auth.py
@@ -82,16 +82,16 @@ class AuthTestCase(unittest.TestCase):
         self.hs.clock.now = 1000
 
         token = self.macaroon_generator.generate_short_term_login_token("a_user", 5000)
-        user_id = yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
-            token
+        user_id = yield defer.ensureDeferred(
+            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):
-            yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
-                token
+            yield defer.ensureDeferred(
+                self.auth_handler.validate_short_term_login_token_and_get_user_id(token)
             )
 
     @defer.inlineCallbacks
@@ -99,8 +99,10 @@ class AuthTestCase(unittest.TestCase):
         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()
+        user_id = yield defer.ensureDeferred(
+            self.auth_handler.validate_short_term_login_token_and_get_user_id(
+                macaroon.serialize()
+            )
         )
         self.assertEqual("a_user", user_id)
 
@@ -109,20 +111,26 @@ class AuthTestCase(unittest.TestCase):
         macaroon.add_first_party_caveat("user_id = b_user")
 
         with self.assertRaises(synapse.api.errors.AuthError):
-            yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
-                macaroon.serialize()
+            yield defer.ensureDeferred(
+                self.auth_handler.validate_short_term_login_token_and_get_user_id(
+                    macaroon.serialize()
+                )
             )
 
     @defer.inlineCallbacks
     def test_mau_limits_disabled(self):
         self.hs.config.limit_usage_by_mau = False
         # Ensure does not throw exception
-        yield self.auth_handler.get_access_token_for_user_id(
-            "user_a", device_id=None, valid_until_ms=None
+        yield defer.ensureDeferred(
+            self.auth_handler.get_access_token_for_user_id(
+                "user_a", device_id=None, valid_until_ms=None
+            )
         )
 
-        yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
-            self._get_macaroon().serialize()
+        yield defer.ensureDeferred(
+            self.auth_handler.validate_short_term_login_token_and_get_user_id(
+                self._get_macaroon().serialize()
+            )
         )
 
     @defer.inlineCallbacks
@@ -133,16 +141,20 @@ class AuthTestCase(unittest.TestCase):
         )
 
         with self.assertRaises(ResourceLimitError):
-            yield self.auth_handler.get_access_token_for_user_id(
-                "user_a", device_id=None, valid_until_ms=None
+            yield defer.ensureDeferred(
+                self.auth_handler.get_access_token_for_user_id(
+                    "user_a", device_id=None, valid_until_ms=None
+                )
             )
 
         self.hs.get_datastore().get_monthly_active_count = Mock(
             return_value=defer.succeed(self.large_number_of_users)
         )
         with self.assertRaises(ResourceLimitError):
-            yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
-                self._get_macaroon().serialize()
+            yield defer.ensureDeferred(
+                self.auth_handler.validate_short_term_login_token_and_get_user_id(
+                    self._get_macaroon().serialize()
+                )
             )
 
     @defer.inlineCallbacks
@@ -154,16 +166,20 @@ class AuthTestCase(unittest.TestCase):
             return_value=defer.succeed(self.hs.config.max_mau_value)
         )
         with self.assertRaises(ResourceLimitError):
-            yield self.auth_handler.get_access_token_for_user_id(
-                "user_a", device_id=None, valid_until_ms=None
+            yield defer.ensureDeferred(
+                self.auth_handler.get_access_token_for_user_id(
+                    "user_a", device_id=None, valid_until_ms=None
+                )
             )
 
         self.hs.get_datastore().get_monthly_active_count = Mock(
             return_value=defer.succeed(self.hs.config.max_mau_value)
         )
         with self.assertRaises(ResourceLimitError):
-            yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
-                self._get_macaroon().serialize()
+            yield defer.ensureDeferred(
+                self.auth_handler.validate_short_term_login_token_and_get_user_id(
+                    self._get_macaroon().serialize()
+                )
             )
         # If in monthly active cohort
         self.hs.get_datastore().user_last_seen_monthly_active = Mock(
@@ -172,8 +188,10 @@ class AuthTestCase(unittest.TestCase):
         self.hs.get_datastore().get_monthly_active_count = Mock(
             return_value=defer.succeed(self.hs.config.max_mau_value)
         )
-        yield self.auth_handler.get_access_token_for_user_id(
-            "user_a", device_id=None, valid_until_ms=None
+        yield defer.ensureDeferred(
+            self.auth_handler.get_access_token_for_user_id(
+                "user_a", device_id=None, valid_until_ms=None
+            )
         )
         self.hs.get_datastore().user_last_seen_monthly_active = Mock(
             return_value=defer.succeed(self.hs.get_clock().time_msec())
@@ -181,8 +199,10 @@ class AuthTestCase(unittest.TestCase):
         self.hs.get_datastore().get_monthly_active_count = Mock(
             return_value=defer.succeed(self.hs.config.max_mau_value)
         )
-        yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
-            self._get_macaroon().serialize()
+        yield defer.ensureDeferred(
+            self.auth_handler.validate_short_term_login_token_and_get_user_id(
+                self._get_macaroon().serialize()
+            )
         )
 
     @defer.inlineCallbacks
@@ -193,15 +213,19 @@ class AuthTestCase(unittest.TestCase):
             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", device_id=None, valid_until_ms=None
+        yield defer.ensureDeferred(
+            self.auth_handler.get_access_token_for_user_id(
+                "user_a", device_id=None, valid_until_ms=None
+            )
         )
 
         self.hs.get_datastore().get_monthly_active_count = Mock(
             return_value=defer.succeed(self.small_number_of_users)
         )
-        yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
-            self._get_macaroon().serialize()
+        yield defer.ensureDeferred(
+            self.auth_handler.validate_short_term_login_token_and_get_user_id(
+                self._get_macaroon().serialize()
+            )
         )
 
     def _get_macaroon(self):
diff --git a/tests/handlers/test_register.py b/tests/handlers/test_register.py
index e7b638dbfe..f1dc51d6c9 100644
--- a/tests/handlers/test_register.py
+++ b/tests/handlers/test_register.py
@@ -294,7 +294,9 @@ class RegistrationTestCase(unittest.HomeserverTestCase):
                 create_profile_with_displayname=user.localpart,
             )
         else:
-            yield self.hs.get_auth_handler().delete_access_tokens_for_user(user_id)
+            yield defer.ensureDeferred(
+                self.hs.get_auth_handler().delete_access_tokens_for_user(user_id)
+            )
 
         yield self.store.add_access_token_to_user(
             user_id=user_id, token=token, device_id=None, valid_until_ms=None