1 files changed, 22 insertions, 4 deletions
diff --git a/tests/rest/client/v2_alpha/test_register.py b/tests/rest/client/v2_alpha/test_register.py
index cda0a2b27c..8ac56a1fb2 100644
--- a/tests/rest/client/v2_alpha/test_register.py
+++ b/tests/rest/client/v2_alpha/test_register.py
@@ -30,6 +30,7 @@ class RegisterRestServletTestCase(unittest.TestCase):
         self.registration_handler = Mock()
         self.identity_handler = Mock()
         self.login_handler = Mock()
+        self.device_handler = Mock()
 
         # do the dance to hook it up to the hs global
         self.handlers = Mock(
@@ -42,6 +43,7 @@ class RegisterRestServletTestCase(unittest.TestCase):
         self.hs.get_auth = Mock(return_value=self.auth)
         self.hs.get_handlers = Mock(return_value=self.handlers)
         self.hs.get_auth_handler = Mock(return_value=self.auth_handler)
+        self.hs.get_device_handler = Mock(return_value=self.device_handler)
         self.hs.config.enable_registration = True
 
         # init the thing we're testing
@@ -61,13 +63,18 @@ class RegisterRestServletTestCase(unittest.TestCase):
             "id": "1234"
         }
         self.registration_handler.appservice_register = Mock(
-            return_value=(user_id, token)
+            return_value=user_id
         )
+        self.auth_handler.get_login_tuple_for_user_id = Mock(
+            return_value=(token, "kermits_refresh_token")
+        )
+
         (code, result) = yield self.servlet.on_POST(self.request)
         self.assertEquals(code, 200)
         det_data = {
             "user_id": user_id,
             "access_token": token,
+            "refresh_token": "kermits_refresh_token",
             "home_server": self.hs.hostname
         }
         self.assertDictContainsSubset(det_data, result)
@@ -105,26 +112,37 @@ class RegisterRestServletTestCase(unittest.TestCase):
     def test_POST_user_valid(self):
         user_id = "@kermit:muppet"
         token = "kermits_access_token"
+        device_id = "frogfone"
         self.request_data = json.dumps({
             "username": "kermit",
-            "password": "monkey"
+            "password": "monkey",
+            "device_id": device_id,
         })
         self.registration_handler.check_username = Mock(return_value=True)
         self.auth_result = (True, None, {
             "username": "kermit",
             "password": "monkey"
         }, None)
-        self.registration_handler.register = Mock(return_value=(user_id, token))
+        self.registration_handler.register = Mock(return_value=(user_id, None))
+        self.auth_handler.get_login_tuple_for_user_id = Mock(
+            return_value=(token, "kermits_refresh_token")
+        )
+        self.device_handler.check_device_registered = \
+            Mock(return_value=device_id)
 
         (code, result) = yield self.servlet.on_POST(self.request)
         self.assertEquals(code, 200)
         det_data = {
             "user_id": user_id,
             "access_token": token,
-            "home_server": self.hs.hostname
+            "refresh_token": "kermits_refresh_token",
+            "home_server": self.hs.hostname,
+            "device_id": device_id,
         }
         self.assertDictContainsSubset(det_data, result)
         self.assertIn("refresh_token", result)
+        self.auth_handler.get_login_tuple_for_user_id(
+            user_id, device_id=device_id, initial_device_display_name=None)
 
     def test_POST_disabled_registration(self):
         self.hs.config.enable_registration = False