1 files changed, 13 insertions, 8 deletions
diff --git a/tests/handlers/test_register.py b/tests/handlers/test_register.py
index 4ea59a58de..35d1bcab3e 100644
--- a/tests/handlers/test_register.py
+++ b/tests/handlers/test_register.py
@@ -17,7 +17,7 @@ from mock import Mock
 
 from twisted.internet import defer
 
-from synapse.api.errors import RegistrationError
+from synapse.api.errors import AuthError
 from synapse.handlers.register import RegistrationHandler
 from synapse.types import UserID, create_requester
 
@@ -40,13 +40,15 @@ class RegistrationTestCase(unittest.TestCase):
         self.mock_distributor.declare("registered_user")
         self.mock_captcha_client = Mock()
         self.hs = yield setup_test_homeserver(
+            self.addCleanup,
             handlers=None,
             http_client=None,
             expire_access_token=True,
             profile_handler=Mock(),
         )
         self.macaroon_generator = Mock(
-            generate_access_token=Mock(return_value='secret'))
+            generate_access_token=Mock(return_value='secret')
+        )
         self.hs.get_macaroon_generator = Mock(return_value=self.macaroon_generator)
         self.hs.handlers = RegistrationHandlers(self.hs)
         self.handler = self.hs.get_handlers().registration_handler
@@ -62,7 +64,8 @@ class RegistrationTestCase(unittest.TestCase):
         user_id = "@someone:test"
         requester = create_requester("@as:test")
         result_user_id, result_token = yield self.handler.get_or_create_user(
-            requester, local_part, display_name)
+            requester, local_part, display_name
+        )
         self.assertEquals(result_user_id, user_id)
         self.assertEquals(result_token, 'secret')
 
@@ -73,13 +76,15 @@ class RegistrationTestCase(unittest.TestCase):
         yield store.register(
             user_id=frank.to_string(),
             token="jkv;g498752-43gj['eamb!-5",
-            password_hash=None)
+            password_hash=None,
+        )
         local_part = "frank"
         display_name = "Frank"
         user_id = "@frank:test"
         requester = create_requester("@as:test")
         result_user_id, result_token = yield self.handler.get_or_create_user(
-            requester, local_part, display_name)
+            requester, local_part, display_name
+        )
         self.assertEquals(result_user_id, user_id)
         self.assertEquals(result_token, 'secret')
 
@@ -104,7 +109,7 @@ class RegistrationTestCase(unittest.TestCase):
         self.store.get_monthly_active_count = Mock(
             return_value=defer.succeed(self.lots_of_users)
         )
-        with self.assertRaises(RegistrationError):
+        with self.assertRaises(AuthError):
             yield self.handler.get_or_create_user("requester", 'b', "display_name")
 
     @defer.inlineCallbacks
@@ -113,7 +118,7 @@ class RegistrationTestCase(unittest.TestCase):
         self.store.get_monthly_active_count = Mock(
             return_value=defer.succeed(self.lots_of_users)
         )
-        with self.assertRaises(RegistrationError):
+        with self.assertRaises(AuthError):
             yield self.handler.register(localpart="local_part")
 
     @defer.inlineCallbacks
@@ -122,5 +127,5 @@ class RegistrationTestCase(unittest.TestCase):
         self.store.get_monthly_active_count = Mock(
             return_value=defer.succeed(self.lots_of_users)
         )
-        with self.assertRaises(RegistrationError):
+        with self.assertRaises(AuthError):
             yield self.handler.register_saml2(localpart="local_part")