1 files changed, 9 insertions, 10 deletions
diff --git a/tests/handlers/test_profile.py b/tests/handlers/test_profile.py
index 891c986fbc..5af3db2cd5 100644
--- a/tests/handlers/test_profile.py
+++ b/tests/handlers/test_profile.py
@@ -14,12 +14,13 @@
 # limitations under the License.
 
 
-from mock import Mock, NonCallableMock
+from mock import Mock, patch
 
 from twisted.internet import defer
 
 import synapse.types
 from synapse.api.errors import AuthError, SynapseError
+from synapse.api.ratelimiting import Ratelimiter
 from synapse.handlers.profile import MasterProfileHandler
 from synapse.types import UserID
 
@@ -55,17 +56,15 @@ class ProfileTestCase(unittest.TestCase):
             federation_client=self.mock_federation,
             federation_server=Mock(),
             federation_registry=self.mock_registry,
-            request_ratelimiter=NonCallableMock(
-                spec_set=["can_do_action", "ratelimit"]
-            ),
-            login_ratelimiter=NonCallableMock(spec_set=["can_do_action", "ratelimit"]),
         )
 
-        self.request_ratelimiter = hs.get_request_ratelimiter()
-        self.request_ratelimiter.can_do_action.return_value = (True, 0)
-
-        self.login_ratelimiter = hs.get_login_ratelimiter()
-        self.login_ratelimiter.can_do_action.return_value = (True, 0)
+        # Patch Ratelimiter to allow all requests
+        patch.object(
+            Ratelimiter, "can_do_action", new_callable=lambda *args, **kwargs: (True, 0.0)
+        )
+        patch.object(
+            Ratelimiter, "ratelimit", new_callable=lambda *args, **kwargs: None
+        )
 
         self.store = hs.get_datastore()