Replace make_awaitable with AsyncMock (#16179)

Python 3.8 provides a native AsyncMock, we can replace the
homegrown version we have.

1 files changed, 13 insertions, 14 deletions

diff --git a/tests/handlers/test_auth.py b/tests/handlers/test_auth.py

index 036dbbc45b..413ff8795b 100644
--- a/tests/handlers/test_auth.py
+++ b/tests/handlers/test_auth.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from typing import Optional
-from unittest.mock import Mock
+from unittest.mock import AsyncMock
 
 import pymacaroons
 
@@ -25,7 +25,6 @@ from synapse.server import HomeServer
 from synapse.util import Clock
 
 from tests import unittest
-from tests.test_utils import make_awaitable
 
 
 class AuthTestCase(unittest.HomeserverTestCase):
@@ -166,8 +165,8 @@ class AuthTestCase(unittest.HomeserverTestCase):
 
     def test_mau_limits_exceeded_large(self) -> None:
         self.auth_blocking._limit_usage_by_mau = True
-        self.hs.get_datastores().main.get_monthly_active_count = Mock(
-            return_value=make_awaitable(self.large_number_of_users)
+        self.hs.get_datastores().main.get_monthly_active_count = AsyncMock(
+            return_value=self.large_number_of_users
         )
 
         self.get_failure(
@@ -177,8 +176,8 @@ class AuthTestCase(unittest.HomeserverTestCase):
             ResourceLimitError,
         )
 
-        self.hs.get_datastores().main.get_monthly_active_count = Mock(
-            return_value=make_awaitable(self.large_number_of_users)
+        self.hs.get_datastores().main.get_monthly_active_count = AsyncMock(
+            return_value=self.large_number_of_users
         )
         token = self.get_success(
             self.auth_handler.create_login_token_for_user_id(self.user1)
@@ -191,8 +190,8 @@ class AuthTestCase(unittest.HomeserverTestCase):
         self.auth_blocking._limit_usage_by_mau = True
 
         # Set the server to be at the edge of too many users.
-        self.hs.get_datastores().main.get_monthly_active_count = Mock(
-            return_value=make_awaitable(self.auth_blocking._max_mau_value)
+        self.hs.get_datastores().main.get_monthly_active_count = AsyncMock(
+            return_value=self.auth_blocking._max_mau_value
         )
 
         # If not in monthly active cohort
@@ -208,8 +207,8 @@ class AuthTestCase(unittest.HomeserverTestCase):
         self.assertIsNone(self.token_login(token))
 
         # If in monthly active cohort
-        self.hs.get_datastores().main.user_last_seen_monthly_active = Mock(
-            return_value=make_awaitable(self.clock.time_msec())
+        self.hs.get_datastores().main.user_last_seen_monthly_active = AsyncMock(
+            return_value=self.clock.time_msec()
         )
         self.get_success(
             self.auth_handler.create_access_token_for_user_id(
@@ -224,8 +223,8 @@ class AuthTestCase(unittest.HomeserverTestCase):
     def test_mau_limits_not_exceeded(self) -> None:
         self.auth_blocking._limit_usage_by_mau = True
 
-        self.hs.get_datastores().main.get_monthly_active_count = Mock(
-            return_value=make_awaitable(self.small_number_of_users)
+        self.hs.get_datastores().main.get_monthly_active_count = AsyncMock(
+            return_value=self.small_number_of_users
         )
         # Ensure does not raise exception
         self.get_success(
@@ -234,8 +233,8 @@ class AuthTestCase(unittest.HomeserverTestCase):
             )
         )
 
-        self.hs.get_datastores().main.get_monthly_active_count = Mock(
-            return_value=make_awaitable(self.small_number_of_users)
+        self.hs.get_datastores().main.get_monthly_active_count = AsyncMock(
+            return_value=self.small_number_of_users
         )
         token = self.get_success(
             self.auth_handler.create_login_token_for_user_id(self.user1)