check that user model has roles attr on user create

flask_security/datastore.py

@@ -126,11 +126,26 @@ def _prepare_role_modify_args(self, user, role):
     def _prepare_create_user_args(self, **kwargs):
         kwargs.setdefault('active', True)
         roles = kwargs.get('roles', [])
-        for i, role in enumerate(roles):
-            rn = role.name if isinstance(role, self.role_model) else role
-            # see if the role exists
-            roles[i] = self.find_role(rn)
-        kwargs['roles'] = roles
+
+        if not hasattr(self.user_model, 'roles') and roles:
+            raise ValueError((
+                'Cannot create user with roles %s, because the user '
+                'model has no "roles" attribute') % str(roles))
+
+        if roles:
+            roles_to_add = []
+
+            for role in roles:
+                rn = role.name if isinstance(role, self.role_model) else role
+
+                # see if the role exists
+                role_to_add = self.find_role(rn)
+
+                if role_to_add:
+                    roles_to_add.append(role_to_add)
+
+            kwargs['roles'] = roles_to_add
+
         return kwargs
 
     def get_user(self, id_or_email):

tests/test_datastore.py

@@ -21,6 +21,19 @@ class Role(RoleMixin):
     pass
 
 
+class FakeUserDatastore(UserDatastore):
+    def find_role(self, role_name):
+        return role_name if role_name != 'lordofchaos' else None
+
+
+class FakeUserModel(object):
+    roles = 'ABC'
+
+
+class FakeUserModelWithoutRoles(object):
+    pass
+
+
 def test_unimplemented_datastore_methods():
     datastore = Datastore(None)
     assert datastore.db is None
@@ -183,3 +196,28 @@ def test_init_app_kwargs_override_constructor_kwargs(app, datastore):
 
     assert security.login_form == 'init_app_login_form'
     assert security.register_form == '__init__register_form'
+
+
+def test_prepare_create_user_args_with_roles():
+    datastore = FakeUserDatastore(FakeUserModel(), type(None))
+    roles = ['godofthunder', 'lordofchaos', 'majordufus']
+
+    create_kwargs = datastore._prepare_create_user_args(roles=roles)
+
+    assert create_kwargs == {
+        'active': True, 'roles': ['godofthunder', 'majordufus']}
+
+
+def test_prepare_create_user_args_without_roles():
+    datastore = FakeUserDatastore(FakeUserModel(), type(None))
+
+    create_kwargs = datastore._prepare_create_user_args(active=False)
+
+    assert create_kwargs == {'active': False}
+
+
+def test_prepare_create_user_args_without_usermodel_roles():
+    datastore = FakeUserDatastore(FakeUserModelWithoutRoles(), type(None))
+
+    with raises(ValueError):
+        datastore._prepare_create_user_args(roles=['hippotamer'])