diff --git a/flask_security/datastore.py b/flask_security/datastore.py index 3896b341..22786c0d 100644 --- a/flask_security/datastore.py +++ b/flask_security/datastore.py @@ -125,14 +125,26 @@ def _prepare_role_modify_args(self, user, role): def _prepare_create_user_args(self, **kwargs): kwargs.setdefault('active', True) + roles = kwargs.get('roles', []) - if hasattr(self.user_model, 'roles'): - roles = kwargs.get('roles', []) - for i, role in enumerate(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 - roles[i] = self.find_role(rn) - kwargs['roles'] = roles + 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 diff --git a/tests/test_datastore.py b/tests/test_datastore.py index 34a7735f..2ea3fbec 100644 --- a/tests/test_datastore.py +++ b/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'])