From 941df0366d60abbf186f40385af9326b885da01f Mon Sep 17 00:00:00 2001
From: Marti Raudsepp <marti@juffo.org>
Date: Thu, 17 Aug 2017 19:24:10 +0300
Subject: [PATCH] Fix roles display on user screen and fix removing user roles
 (#879)

---
 lemur/tests/test_users.py | 34 ++++++++++++++++++++++++++++++++++
 lemur/users/schemas.py    |  1 +
 lemur/users/service.py    |  4 ++--
 3 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/lemur/tests/test_users.py b/lemur/tests/test_users.py
index e71980a3..e81fbbd2 100644
--- a/lemur/tests/test_users.py
+++ b/lemur/tests/test_users.py
@@ -1,5 +1,8 @@
+import json
+
 import pytest
 
+from lemur.tests.factories import UserFactory, RoleFactory
 from lemur.users.views import *  # noqa
 from .vectors import VALID_ADMIN_HEADER_TOKEN, VALID_USER_HEADER_TOKEN
 
@@ -107,3 +110,34 @@ def test_sensitive_filter(client):
 def test_sensitive_sort(client):
     resp = client.get(api.url_for(UsersList) + '?sortBy=password&sortDir=asc', headers=VALID_ADMIN_HEADER_TOKEN)
     assert "'password' is not sortable or filterable" in resp.json['message']
+
+
+def test_user_role_changes(client, session):
+    user = UserFactory()
+    role1 = RoleFactory()
+    role2 = RoleFactory()
+    session.flush()
+
+    data = {
+        'active': True,
+        'id': user.id,
+        'username': user.username,
+        'email': user.email,
+        'roles': [
+            {'id': role1.id},
+            {'id': role2.id},
+        ],
+    }
+
+    # PUT two roles
+    resp = client.put(api.url_for(Users, user_id=user.id), data=json.dumps(data), headers=VALID_ADMIN_HEADER_TOKEN)
+    assert resp.status_code == 200
+    assert len(resp.json['roles']) == 2
+    assert set(user.roles) == {role1, role2}
+
+    # Remove one role and PUT again
+    del data['roles'][1]
+    resp = client.put(api.url_for(Users, user_id=user.id), data=json.dumps(data), headers=VALID_ADMIN_HEADER_TOKEN)
+    assert resp.status_code == 200
+    assert len(resp.json['roles']) == 1
+    assert set(user.roles) == {role1}
diff --git a/lemur/users/schemas.py b/lemur/users/schemas.py
index 4d0078fd..a944a7cf 100644
--- a/lemur/users/schemas.py
+++ b/lemur/users/schemas.py
@@ -26,6 +26,7 @@ class UserOutputSchema(LemurOutputSchema):
     username = fields.String()
     email = fields.Email()
     active = fields.Boolean()
+    roles = fields.Nested(AssociatedRoleSchema, many=True)
     profile_picture = fields.String()
 
 
diff --git a/lemur/users/service.py b/lemur/users/service.py
index 40b67843..1b191d7e 100644
--- a/lemur/users/service.py
+++ b/lemur/users/service.py
@@ -64,12 +64,12 @@ def update_roles(user, roles):
     :param user:
     :param roles:
     """
-    for ur in roles:
+    for ur in user.roles:
         for r in roles:
             if r.id == ur.id:
                 break
         else:
-            user.roles.remove(r)
+            user.roles.remove(ur)
 
     for r in roles:
         for ur in user.roles: