Small refactor of the fix for consistency

synapse/module_api/__init__.py

@@ -110,7 +110,7 @@ def register(self, localpart, displayname=None, emails=[]):
             "Using deprecated ModuleApi.register which creates a dummy user device."
         )
         user_id = yield self.register_user(localpart, displayname, emails)
-        _, access_token = yield defer.ensureDeferred(self.register_device(user_id))
+        _, access_token = yield self.register_device(user_id)
         return user_id, access_token
 
     def register_user(self, localpart, displayname=None, emails=[]):
@@ -126,7 +126,7 @@ def register_user(self, localpart, displayname=None, emails=[]):
                 'errcode' property for more information on the reason for failure
 
         Returns:
-            Deferred[str]: user_id
+            defer.Deferred[str]: user_id
         """
         return defer.ensureDeferred(
             self._hs.get_registration_handler().register_user(
@@ -149,10 +149,12 @@ def register_device(self, user_id, device_id=None, initial_display_name=None):
         Returns:
             defer.Deferred[tuple[str, str]]: Tuple of device ID and access token
         """
-        return self._hs.get_registration_handler().register_device(
-            user_id=user_id,
-            device_id=device_id,
-            initial_display_name=initial_display_name,
+        return defer.ensureDeferred(
+            self._hs.get_registration_handler().register_device(
+                user_id=user_id,
+                device_id=device_id,
+                initial_display_name=initial_display_name,
+            )
         )
 
     def record_user_external_id(