Fixed bug in /u/ where home project group membership was lost after edit.

Rather than understanding the code, I rewrote the editing and added a
unit test for it.

pillar/web/users/routes.py

@@ -210,34 +210,7 @@ def users_edit(user_id):
 
     form = UserEditForm()
     if form.validate_on_submit():
-        def get_groups(roles):
-            """Return a set of role ids matching the group names provided"""
-            groups_set = set()
-            for system_role in roles:
-                group = Group.find_one({'where': "name=='%s'" % system_role}, api=api)
-                groups_set.add(group._id)
-            return groups_set
-
-        # Remove any of the default roles
-        system_roles = set([role[0] for role in form.roles.choices])
-        system_groups = get_groups(system_roles)
-        # Current user roles
-        user_roles_list = user.roles if user.roles else []
-        user_roles = set(user_roles_list)
-        user_groups = get_groups(user_roles_list)
-        # Remove all form roles from current roles
-        user_roles = list(user_roles.difference(system_roles))
-        user_groups = list(user_groups.difference(system_groups))
-        # Get the assigned roles
-        system_roles_assigned = form.roles.data
-        system_groups_assigned = get_groups(system_roles_assigned)
-        # Reassign roles based on form.roles.data by adding them to existing roles
-        user_roles += system_roles_assigned
-        user_groups += list(get_groups(user_roles))
-        # Fetch the group for the assigned system roles
-        user.roles = user_roles
-        user.groups = user_groups
-        user.update(api=api)
+        _users_edit(form, user, api)
     else:
         form.roles.data = user.roles
     return render_template('users/edit_embed.html',
@@ -245,6 +218,29 @@ def users_edit(user_id):
         form=form)
 
 
+def _users_edit(form, user, api):
+    """Performs the actual user editing."""
+
+    from pillar.api.service import role_to_group_id, ROLES_WITH_GROUPS
+
+    current_user_roles = set(user.roles or [])
+    current_user_groups = set(user.groups or [])
+
+    roles_in_form = set(form.roles.data)
+
+    granted_roles = roles_in_form - current_user_roles
+    revoked_roles = ROLES_WITH_GROUPS - roles_in_form
+
+    # role_to_group_id contains ObjectIDs, but the SDK works with strings.
+    granted_groups = {str(role_to_group_id[role]) for role in granted_roles}
+    revoked_groups = {str(role_to_group_id[role]) for role in revoked_roles}
+
+    user.roles = list((current_user_roles - revoked_roles).union(granted_roles))
+    user.groups = list((current_user_groups - revoked_groups).union(granted_groups))
+
+    user.update(api=api)
+
+
 @blueprint.route('/u')
 @login_required
 def users_index():

tests/test_web/test_user_admin.py

@@ -0,0 +1,80 @@
+# -*- encoding: utf-8 -*-
+
+"""Unit tests for the user admin interface."""
+
+from __future__ import absolute_import
+
+import json
+import logging
+
+import responses
+from bson import ObjectId
+import flask_login
+import pillarsdk
+from pillar.tests import AbstractPillarTest, TEST_EMAIL_ADDRESS
+from werkzeug import exceptions as wz_exceptions
+
+log = logging.getLogger(__name__)
+
+
+class UserAdminTest(AbstractPillarTest):
+    def setUp(self, **kwargs):
+        AbstractPillarTest.setUp(self, **kwargs)
+        self.create_standard_groups()
+
+        from pillar.api.service import role_to_group_id
+        self.subscriber_gid = role_to_group_id['subscriber']
+        self.demo_gid = role_to_group_id['demo']
+
+    def test_grant_subscriber_role(self):
+        """There was a bug that group membership was lost when a user was edited."""
+
+        import pillar.web.users.routes
+        import pillar.auth
+
+        user_id = self.create_user(roles=())
+        self.create_valid_auth_token(user_id, 'token')
+
+        # Try to access the home project, creating it.
+        self.get('/api/bcloud/home-project',
+                 auth_token='token')
+
+        def get_dbuser():
+            with self.app.test_request_context():
+                db = self.app.db()
+                return db['users'].find_one({'_id': user_id})
+
+        db_user = get_dbuser()
+        groups_pre_op = db_user['groups']
+        self.assertEqual(1, len(groups_pre_op))
+        home_project_gid = groups_pre_op[0]
+
+        # Edit the user, granting subscriber and demo roles.
+        admin_id = 24 * 'a'
+        self.create_user(admin_id, roles=['admin'])
+        self.create_valid_auth_token(admin_id, 'admin-token')
+
+        def edit_user(roles):
+            from pillar.web import system_util
+            from pillar.web.users.forms import UserEditForm
+
+            with self.app.test_request_context():
+                api = system_util.pillar_api(token='admin-token')
+                user = pillarsdk.User.find(user_id, api=api)
+
+                form = UserEditForm()
+                form.roles.data = roles
+                pillar.web.users.routes._users_edit(form, user, api)
+
+        edit_user(['subscriber', 'demo'])
+
+        # Re-check the user group membership.
+        groups = get_dbuser()['groups']
+        self.assertEqual({home_project_gid, self.subscriber_gid, self.demo_gid},
+                         set(groups))
+
+        # Edit user again, revoking demo role.
+        edit_user(['subscriber'])
+        groups = get_dbuser()['groups']
+        self.assertEqual({home_project_gid, self.subscriber_gid},
+                         set(groups))