diff --git a/tests/rest/client/v2_alpha/test_register.py b/tests/rest/client/v2_alpha/test_register.py index 753d5c3e80..18080ebfd6 100644 --- a/tests/rest/client/v2_alpha/test_register.py +++ b/tests/rest/client/v2_alpha/test_register.py @@ -32,7 +32,18 @@ class RegisterRestServletTestCase(unittest.HomeserverTestCase): self.identity_handler = Mock() self.login_handler = Mock() self.device_handler = Mock() - self.device_handler.check_device_registered = Mock(return_value="FAKE") + + def check_device_registered(user_id, device_id, initial_display_name): + # Just echo back the given device ID, or return a new "FAKE" device + # ID + if device_id: + return device_id + else: + return "FAKE" + + self.device_handler.check_device_registered = Mock( + side_effect=check_device_registered, + ) self.datastore = Mock(return_value=Mock()) self.datastore.get_current_state_deltas = Mock(return_value=[]) @@ -106,14 +117,12 @@ class RegisterRestServletTestCase(unittest.HomeserverTestCase): user_id = "@kermit:muppet" token = "kermits_access_token" device_id = "frogfone" - request_data = json.dumps( - {"username": "kermit", "password": "monkey", "device_id": device_id} - ) + params = {"username": "kermit", "password": "monkey", "device_id": device_id} + request_data = json.dumps(params) self.registration_handler.check_username = Mock(return_value=True) - self.auth_result = (None, {"username": "kermit", "password": "monkey"}, None) + self.auth_result = (None, params, None) self.registration_handler.register = Mock(return_value=(user_id, None)) self.auth_handler.get_access_token_for_user_id = Mock(return_value=token) - self.device_handler.check_device_registered = Mock(return_value=device_id) request, channel = self.make_request(b"POST", self.url, request_data) self.render(request)