Address comments on code review

endpoints/api/user.py

@@ -12,6 +12,8 @@ from endpoints.api import (ApiResource, nickname, resource, validate_json_reques
                            license_error)
 from endpoints.api.subscribe import subscribe
 from endpoints.common import common_login
+from endpoints.api.team import try_accept_invite
+
 from data import model
 from data.billing import get_plan
 from auth.permissions import (AdministerOrganizationPermission, CreateRepositoryPermission,
@@ -20,6 +22,7 @@ from auth.auth_context import get_authenticated_user
 from auth import scopes
 from util.gravatar import compute_hash
 from util.useremails import (send_confirmation_email, send_recovery_email, send_change_email)
+from util.names import parse_single_urn
 
 import features
 
@@ -179,11 +182,15 @@ class User(ApiResource):
     return user_view(user)
 
   @nickname('createNewUser')
+  @parse_args
+  @query_param('inviteCode', 'Invitation code given for creating the user.', type=str,
+               default='')
   @internal_only
   @validate_json_request('NewUser')
-  def post(self):
+  def post(self, args):
     """ Create a new user. """
     user_data = request.get_json()
+    invite_code = args['inviteCode']
 
     existing_user = model.get_user(user_data['username'])
     if existing_user:
@@ -194,6 +201,14 @@ class User(ApiResource):
                                    user_data['email'])
       code = model.create_confirm_email_code(new_user)
       send_confirmation_email(new_user.username, new_user.email, code.code)
+
+      # Handle any invite codes.
+      parsed_invite = parse_single_urn(invite_code)
+      if parsed_invite is not None:
+        if parsed_invite[0] == 'teaminvite':
+          # Add the user to the team.
+          try_accept_invite(invite_code, new_user)
+
       return 'Created', 201
     except model.TooManyUsersException as ex:
       raise license_error(exception=ex)