Refactored role methods

jupyterhub/apihandlers/users.py

@@ -14,6 +14,7 @@ from tornado import web
 from tornado.iostream import StreamClosedError
 
 from .. import orm
+from ..roles import assign_default_roles
 from ..roles import update_roles
 from ..scopes import needs_scope
 from ..user import User
@@ -151,7 +152,7 @@ class UserListAPIHandler(APIHandler):
             user = self.user_from_username(name)
             if admin:
                 user.admin = True
-            update_roles(self.db, obj=user, kind='users')
+            assign_default_roles(self.db, entity=user)
             self.db.commit()
             try:
                 await maybe_future(self.authenticator.add_user(user))
@@ -218,7 +219,7 @@ class UserAPIHandler(APIHandler):
             self._check_user_model(data)
             if 'admin' in data:
                 user.admin = data['admin']
-                update_roles(self.db, obj=user, kind='users')
+                assign_default_roles(self.db, entity=user)
         self.db.commit()
 
         try:
@@ -286,7 +287,7 @@ class UserAPIHandler(APIHandler):
             else:
                 setattr(user, key, value)
                 if key == 'admin':
-                    update_roles(self.db, obj=user, kind='users')
+                    assign_default_roles(self.db, entity=user)
         self.db.commit()
         user_ = self.user_model(user)
         user_['auth_state'] = await user.get_auth_state()

jupyterhub/app.py

@@ -1857,15 +1857,16 @@ class JupyterHub(Application):
         # load default roles
         default_roles = roles.get_default_roles()
         for role in default_roles:
-            roles.add_role(db, role)
+            roles.create_role(db, role)
 
         # load predefined roles from config file
         for predef_role in self.load_roles:
-            roles.add_role(db, predef_role)
+            roles.create_role(db, predef_role)
             # add users, services and/or tokens
             for bearer in role_bearers:
                 if bearer in predef_role.keys():
                     for bname in predef_role[bearer]:
+
                         if bearer == 'users':
                             bname = self.authenticator.normalize_username(bname)
                             if not (
@@ -1877,8 +1878,10 @@ class JupyterHub(Application):
                                     "Username %r is not in Authenticator.allowed_users"
                                     % bname
                                 )
-                        roles.add_obj(
+                        Class = orm.get_class(bearer)
-                            db, objname=bname, kind=bearer, rolename=predef_role['name']
+                        orm_obj = Class.find(db, bname)
+                        roles.grant_role(
+                            db, entity=orm_obj, rolename=predef_role['name']
                         )
 
         # make sure all users, services and tokens have at least one role (update with default)
@@ -1886,7 +1889,7 @@ class JupyterHub(Application):
             Class = orm.get_class(bearer)
             for obj in db.query(Class):
                 if len(obj.roles) < 1:
-                    roles.update_roles(db, obj=obj, kind=bearer)
+                    roles.assign_default_roles(db, entity=obj)
         db.commit()
 
     async def _add_tokens(self, token_dict, kind):

jupyterhub/handlers/base.py

@@ -480,7 +480,7 @@ class BaseHandler(RequestHandler):
             # not found, create and register user
             u = orm.User(name=username)
             self.db.add(u)
-            roles.update_roles(self.db, obj=u, kind='users')
+            roles.assign_default_roles(self.db, entity=u)
             TOTAL_USERS.inc()
             self.db.commit()
             user = self._user_from_orm(u)
@@ -765,7 +765,7 @@ class BaseHandler(RequestHandler):
         # Only set `admin` if the authenticator returned an explicit value.
         if admin is not None and admin != user.admin:
             user.admin = admin
-            roles.update_roles(self.db, obj=user, kind='users')
+            roles.assign_default_roles(self.db, entity=user)
             self.db.commit()
         # always set auth_state and commit,
         # because there could be key-rotation or clearing of previous values

jupyterhub/orm.py

@@ -39,7 +39,8 @@ from sqlalchemy.types import Text
 from sqlalchemy.types import TypeDecorator
 from tornado.log import app_log
 
-from .roles import add_role
+from .roles import assign_default_roles
+from .roles import create_role
 from .roles import get_default_roles
 from .roles import update_roles
 from .utils import compare_token
@@ -621,8 +622,12 @@ class APIToken(Hashed, Base):
         if not token_role:
             default_roles = get_default_roles()
             for role in default_roles:
-                add_role(db, role)
+                create_role(db, role)
-        update_roles(db, obj=orm_token, kind='tokens', roles=roles)
+        if roles:
+            update_roles(db, entity=orm_token, roles=roles)
+        else:
+            assign_default_roles(db, entity=orm_token)
+
         db.commit()
         return token
 

jupyterhub/roles.py

@@ -73,7 +73,7 @@ def expand_self_scope(name, read_only=False):
     return {"{}!user={}".format(scope, name) for scope in scope_list}
 
 
-def get_scope_hierarchy():
+def _get_scope_hierarchy():
     """
     Returns a dictionary of scopes:
     scopes.keys() = scopes of highest level and scopes that have their own subscopes
@@ -106,7 +106,7 @@ def get_scope_hierarchy():
 def expand_scope(scopename):
     """Returns a set of all subscopes"""
 
-    scopes = get_scope_hierarchy()
+    scopes = _get_scope_hierarchy()
     subscopes = [scopename]
 
     def expand_subscopes(index):
@@ -133,7 +133,7 @@ def expand_scope(scopename):
 
 def expand_roles_to_scopes(orm_object):
     """Get the scopes listed in the roles of the User/Service/Group/Token"""
-    scopes = get_subscopes(*orm_object.roles)
+    scopes = _get_subscopes(*orm_object.roles)
     if 'self' in scopes:
         scopes.remove('self')
         if isinstance(orm_object, orm.User) or hasattr(orm_object, 'orm_user'):
@@ -141,7 +141,7 @@ def expand_roles_to_scopes(orm_object):
     return scopes
 
 
-def get_subscopes(*args):
+def _get_subscopes(*args):
     """Returns a set of all available subscopes for a specified role or list of roles"""
 
     scope_list = []
@@ -154,7 +154,7 @@ def get_subscopes(*args):
     return scopes
 
 
-def add_role(db, role_dict):
+def create_role(db, role_dict):
     """Adds a new role to database or modifies an existing one"""
 
     if 'name' not in role_dict.keys():
@@ -180,105 +180,103 @@ def add_role(db, role_dict):
 def existing_only(func):
     """Decorator for checking if objects and roles exist"""
 
-    def check_existence(db, objname, kind, rolename):
+    def check_existence(db, entity, rolename):
-
-        Class = orm.get_class(kind)
-        obj = Class.find(db, objname)
         role = orm.Role.find(db, rolename)
-
+        if entity is None:
-        if obj is None:
+            raise ValueError(
-            raise ValueError("%r of kind %r does not exist" % (objname, kind))
+                "%r of kind %r does not exist" % (entity, type(entity).__name__)
+            )
         elif role is None:
             raise ValueError("Role %r does not exist" % rolename)
         else:
-            func(db, obj, kind, role)
+            func(db, entity, role)
 
     return check_existence
 
 
 @existing_only
-def add_obj(db, objname, kind, rolename):
+def grant_role(db, entity, rolename):
     """Adds a role for users, services or tokens"""
-
+    if rolename not in entity.roles:
-    if rolename not in objname.roles:
+        entity.roles.append(rolename)
-        objname.roles.append(rolename)
         db.commit()
 
 
 @existing_only
-def remove_obj(db, objname, kind, rolename):
+def strip_role(db, entity, rolename):
     """Removes a role for users, services or tokens"""
-
+    if rolename in entity.roles:
-    if rolename in objname.roles:
+        entity.roles.remove(rolename)
-        objname.roles.remove(rolename)
         db.commit()
 
 
-def switch_default_role(db, obj, kind, admin):
+def _switch_default_role(db, obj, admin):
     """Switch between default user/service and admin roles for users/services"""
-
     user_role = orm.Role.find(db, 'user')
     admin_role = orm.Role.find(db, 'admin')
 
-    def add_and_remove(db, obj, kind, current_role, new_role):
+    def add_and_remove(db, obj, current_role, new_role):
-
         if current_role in obj.roles:
-            remove_obj(db, objname=obj.name, kind=kind, rolename=current_role.name)
+            strip_role(db, entity=obj, rolename=current_role.name)
         # only add new default role if the user has no other roles
         if len(obj.roles) < 1:
-            add_obj(db, objname=obj.name, kind=kind, rolename=new_role.name)
+            grant_role(db, entity=obj, rolename=new_role.name)
 
     if admin:
-        add_and_remove(db, obj, kind, user_role, admin_role)
+        add_and_remove(db, obj, user_role, admin_role)
     else:
-        add_and_remove(db, obj, kind, admin_role, user_role)
+        add_and_remove(db, obj, admin_role, user_role)
 
 
-def update_roles(db, obj, kind, roles=None):
+def assign_default_roles(db, entity):
+    """Assigns the default roles to an entity:
+    users and services get 'user' role, unless they have admin flag
+    Tokens get 'token' role"""
+    default_token_role = orm.Role.find(db, 'token')
+    # tokens can have only 'token' role as default
+    # assign the default only for tokens
+    if isinstance(entity, orm.APIToken):
+        if not entity.roles and entity.user is not None:
+            default_token_role.tokens.append(entity)
+        db.commit()
+    # users and services can have 'user' or 'admin' roles as default
+    else:
+        # todo: when we deprecate admin flag: replace with role check
+        _switch_default_role(db, entity, entity.admin)
+
+
+def update_roles(db, entity, roles):
     """Updates object's roles if specified,
     assigns default if no roles specified"""
-
+    Class = type(entity)
-    Class = orm.get_class(kind)
-    default_token_role = orm.Role.find(db, 'token')
     standard_permissions = {'all', 'read:all'}
-    if roles:
+    for rolename in roles:
-        for rolename in roles:
-            if Class == orm.APIToken:
-
-                role = orm.Role.find(db, rolename)
-                if role:
-                    # compare the requested role permissions with the owner's permissions (scopes)
-                    token_scopes = get_subscopes(role)
-                    extra_scopes = token_scopes - standard_permissions
-                    # find the owner and their roles
-                    owner = None
-                    if obj.user_id:
-                        owner = db.query(orm.User).get(obj.user_id)
-                    elif obj.service_id:
-                        owner = db.query(orm.Service).get(obj.service_id)
-                    if owner:
-                        owner_scopes = expand_roles_to_scopes(owner)
-                        if (extra_scopes).issubset(owner_scopes):
-                            role.tokens.append(obj)
-                        else:
-                            raise ValueError(
-                                'Requested token role %r has more permissions than the token owner: [%s]'
-                                % (rolename, ",".join(extra_scopes - owner_scopes))
-                            )
-                else:
-                    raise NameError('Role %r does not exist' % rolename)
-            else:
-                add_obj(db, objname=obj.name, kind=kind, rolename=rolename)
-    else:
-        # tokens can have only 'token' role as default
-        # assign the default only for tokens
         if Class == orm.APIToken:
-            if not obj.roles and obj.user is not None:
+
-                default_token_role.tokens.append(obj)
+            role = orm.Role.find(db, rolename)
-            db.commit()
+            if role:
-        # users and services can have 'user' or 'admin' roles as default
+                # compare the requested role permissions with the owner's permissions (scopes)
+                token_scopes = _get_subscopes(role)
+                extra_scopes = token_scopes - standard_permissions
+                # find the owner and their roles
+                owner = None
+                if entity.user_id:
+                    owner = db.query(orm.User).get(entity.user_id)
+                elif entity.service_id:
+                    owner = db.query(orm.Service).get(entity.service_id)
+                if owner:
+                    owner_scopes = expand_roles_to_scopes(owner)
+                    if (extra_scopes).issubset(owner_scopes):
+                        role.tokens.append(entity)
+                    else:
+                        raise ValueError(
+                            'Requested token role %r has more permissions than the token owner: [%s]'
+                            % (rolename, ",".join(extra_scopes - owner_scopes))
+                        )
+            else:
+                raise NameError('Role %r does not exist' % rolename)
         else:
-            switch_default_role(db, obj, kind, obj.admin)
+            grant_role(db, entity=entity, rolename=rolename)
 
 
 def mock_roles(app, name, kind):
@@ -287,5 +285,5 @@ def mock_roles(app, name, kind):
     obj = Class.find(app.db, name=name)
     default_roles = get_default_roles()
     for role in default_roles:
-        add_role(app.db, role)
+        create_role(app.db, role)
-    update_roles(db=app.db, obj=obj, kind=kind)
+    assign_default_roles(db=app.db, entity=obj)

jupyterhub/tests/conftest.py

@@ -251,7 +251,7 @@ def _mockservice(request, app, url=False):
         assert name in app._service_map
         service = app._service_map[name]
         token = service.orm.api_tokens[0]
-        update_roles(app.db, token, 'tokens', roles=['token'])
+        update_roles(app.db, token, roles=['token'])
 
         async def start():
             # wait for proxy to be updated before starting the service

jupyterhub/tests/mocking.py

@@ -342,7 +342,7 @@ class MockHub(JupyterHub):
             self.db.add(user)
             self.db.commit()
             metrics.TOTAL_USERS.inc()
-        roles.update_roles(self.db, obj=user, kind='users')
+        roles.assign_default_roles(self.db, entity=user)
         self.db.commit()
 
     def stop(self):

jupyterhub/tests/test_app.py

@@ -50,7 +50,7 @@ def test_raise_error_on_missing_specified_config():
     process = Popen(
         [sys.executable, '-m', 'jupyterhub', '--config', 'not-available.py']
     )
-    # wait inpatiently for the process to exit like we want it to
+    # wait impatiently for the process to exit like we want it to
     for i in range(100):
         time.sleep(0.1)
         returncode = process.poll()

jupyterhub/tests/test_roles.py

@@ -180,9 +180,9 @@ def test_orm_roles_delete_cascade(db):
 )
 def test_get_subscopes(db, scopes, subscopes):
     """Test role scopes expansion into their subscopes"""
-    roles.add_role(db, {'name': 'testing_scopes', 'scopes': scopes})
+    roles.create_role(db, {'name': 'testing_scopes', 'scopes': scopes})
     role = orm.Role.find(db, name='testing_scopes')
-    response = roles.get_subscopes(role)
+    response = roles._get_subscopes(role)
     assert response == subscopes
     db.delete(role)
 
@@ -386,8 +386,8 @@ async def test_load_roles_tokens(tmpdir, request):
 )
 async def test_get_new_token_via_api(app, headers, role_list, status):
     user = add_user(app.db, app, name='user')
-    roles.add_role(app.db, {'name': 'reader', 'scopes': ['all']})
+    roles.create_role(app.db, {'name': 'reader', 'scopes': ['all']})
-    roles.add_role(app.db, {'name': 'user_creator', 'scopes': ['admin:users']})
+    roles.create_role(app.db, {'name': 'user_creator', 'scopes': ['admin:users']})
     if role_list:
         body = json.dumps({'roles': role_list})
     else:

jupyterhub/tests/test_scopes.py

@@ -230,7 +230,7 @@ async def test_expand_groups(app, user_name, in_group, status_code):
             'read:groups',
         ],
     }
-    roles.add_role(app.db, test_role)
+    roles.create_role(app.db, test_role)
     user = add_user(app.db, name=user_name)
     group_name = 'bluth'
     group = orm.Group.find(app.db, name=group_name)
@@ -240,8 +240,8 @@ async def test_expand_groups(app, user_name, in_group, status_code):
     if in_group and user not in group.users:
         group.users.append(user)
     kind = 'users'
-    roles.update_roles(app.db, user, kind, roles=['test'])
+    roles.update_roles(app.db, user, roles=['test'])
-    roles.remove_obj(app.db, user_name, kind, 'user')
+    roles.strip_role(app.db, user, 'user')
     app.db.commit()
     r = await api_request(
         app, 'users', user_name, headers=auth_header(app.db, user_name)
@@ -265,10 +265,10 @@ err_message = "No access to resources or resources not found"
 async def test_request_fake_user(app):
     user_name = 'buster'
     fake_user = 'annyong'
-    add_user(app.db, name=user_name)
+    user = add_user(app.db, name=user_name)
     test_role = generate_test_role(user_name, ['read:users!group=stuff'])
-    roles.add_role(app.db, test_role)
+    roles.create_role(app.db, test_role)
-    roles.add_obj(app.db, objname=user_name, kind='users', rolename='test')
+    roles.grant_role(app.db, entity=user, rolename='test')
     app.db.commit()
     r = await api_request(
         app, 'users', fake_user, headers=auth_header(app.db, user_name)
@@ -284,8 +284,8 @@ async def test_refuse_exceeding_token_permissions(app):
     add_user(app.db, name='user')
     api_token = user.new_api_token()
     exceeding_role = generate_test_role(user_name, ['read:users'], 'exceeding_role')
-    roles.add_role(app.db, exceeding_role)
+    roles.create_role(app.db, exceeding_role)
-    roles.add_obj(app.db, objname=api_token, kind='tokens', rolename='exceeding_role')
+    roles.grant_role(app.db, entity=user.api_tokens[0], rolename='exceeding_role')
     app.db.commit()
     headers = {'Authorization': 'token %s' % api_token}
     r = await api_request(app, 'users', headers=headers)
@@ -304,11 +304,11 @@ async def test_exceeding_user_permissions(app):
     subreader_role = generate_test_role(
         user_name, ['read:users:groups'], 'subreader_role'
     )
-    roles.add_role(app.db, reader_role)
+    roles.create_role(app.db, reader_role)
-    roles.add_role(app.db, subreader_role)
+    roles.create_role(app.db, subreader_role)
     app.db.commit()
-    roles.update_roles(app.db, user, kind='users', roles=['reader_role'])
+    roles.update_roles(app.db, user, roles=['reader_role'])
-    roles.update_roles(app.db, orm_api_token, kind='tokens', roles=['subreader_role'])
+    roles.update_roles(app.db, orm_api_token, roles=['subreader_role'])
     orm_api_token.roles.remove(orm.Role.find(app.db, name='token'))
     app.db.commit()
 
@@ -318,7 +318,7 @@ async def test_exceeding_user_permissions(app):
     keys = {key for user in r.json() for key in user.keys()}
     assert 'groups' in keys
     assert 'last_activity' not in keys
-    roles.remove_obj(app.db, user_name, 'users', 'reader_role')
+    roles.strip_role(app.db, user, 'reader_role')
 
 
 async def test_user_service_separation(app, mockservice_url):
@@ -327,13 +327,11 @@ async def test_user_service_separation(app, mockservice_url):
 
     reader_role = generate_test_role(name, ['read:users'], 'reader_role')
     subreader_role = generate_test_role(name, ['read:users:groups'], 'subreader_role')
-    roles.add_role(app.db, reader_role)
+    roles.create_role(app.db, reader_role)
-    roles.add_role(app.db, subreader_role)
+    roles.create_role(app.db, subreader_role)
     app.db.commit()
-    roles.update_roles(app.db, user, kind='users', roles=['subreader_role'])
+    roles.update_roles(app.db, user, roles=['subreader_role'])
-    roles.update_roles(
+    roles.update_roles(app.db, mockservice_url.orm, roles=['reader_role'])
-        app.db, mockservice_url.orm, kind='services', roles=['reader_role']
-    )
     user.roles.remove(orm.Role.find(app.db, name='user'))
     api_token = user.new_api_token()
     app.db.commit()
@@ -348,12 +346,12 @@ async def test_user_service_separation(app, mockservice_url):
 async def test_request_user_outside_group(app):
     user_name = 'buster'
     fake_user = 'hello'
-    add_user(app.db, name=user_name)
+    user = add_user(app.db, name=user_name)
     add_user(app.db, name=fake_user)
     test_role = generate_test_role(user_name, ['read:users!group=stuff'])
-    roles.add_role(app.db, test_role)
+    roles.create_role(app.db, test_role)
-    roles.add_obj(app.db, objname=user_name, kind='users', rolename='test')
+    roles.grant_role(app.db, entity=user, rolename='test')
-    roles.remove_obj(app.db, objname=user_name, kind='users', rolename='user')
+    roles.strip_role(app.db, entity=user, rolename='user')
     app.db.commit()
     r = await api_request(
         app, 'users', fake_user, headers=auth_header(app.db, user_name)
@@ -369,9 +367,9 @@ async def test_user_filter(app):
     app.db.commit()
     scopes = ['read:users!user=lindsay', 'read:users!user=gob', 'read:users!user=oscar']
     test_role = generate_test_role(user, scopes)
-    roles.add_role(app.db, test_role)
+    roles.create_role(app.db, test_role)
-    roles.add_obj(app.db, objname=user_name, kind='users', rolename='test')
+    roles.grant_role(app.db, entity=user, rolename='test')
-    roles.remove_obj(app.db, objname=user_name, kind='users', rolename='user')
+    roles.strip_role(app.db, entity=user, rolename='user')
     name_in_scope = {'lindsay', 'oscar', 'gob'}
     outside_scope = {'maeby', 'marta'}
     group_name = 'bluth'
@@ -402,8 +400,8 @@ async def test_service_filter(app):
     user = add_user(app.db, name=user_name)
     app.db.commit()
     test_role = generate_test_role(user, ['read:services!service=cull_idle'])
-    roles.add_role(app.db, test_role)
+    roles.create_role(app.db, test_role)
-    roles.add_obj(app.db, objname=user_name, kind='users', rolename='test')
+    roles.grant_role(app.db, entity=user, rolename='test')
     r = await api_request(app, 'services', headers=auth_header(app.db, user_name))
     assert r.status_code == 200
     service_names = set(r.json().keys())
@@ -413,12 +411,12 @@ async def test_service_filter(app):
 async def test_user_filter_with_group(app):
     # Move role setup to setup method?
     user_name = 'sally'
-    add_user(app.db, name=user_name)
+    user = add_user(app.db, name=user_name)
     external_user_name = 'britta'
     add_user(app.db, name=external_user_name)
     test_role = generate_test_role(user_name, ['read:users!group=sitwell'])
-    roles.add_role(app.db, test_role)
+    roles.create_role(app.db, test_role)
-    roles.add_obj(app.db, objname=user_name, kind='users', rolename='test')
+    roles.grant_role(app.db, entity=user, rolename='test')
 
     name_set = {'sally', 'stan'}
     group_name = 'sitwell'
@@ -441,11 +439,11 @@ async def test_user_filter_with_group(app):
 
 async def test_group_scope_filter(app):
     user_name = 'rollerblade'
-    add_user(app.db, name=user_name)
+    user = add_user(app.db, name=user_name)
     scopes = ['read:groups!group=sitwell', 'read:groups!group=bluth']
     test_role = generate_test_role(user_name, scopes)
-    roles.add_role(app.db, test_role)
+    roles.create_role(app.db, test_role)
-    roles.add_obj(app.db, objname=user_name, kind='users', rolename='test')
+    roles.grant_role(app.db, entity=user, rolename='test')
 
     group_set = {'sitwell', 'bluth', 'austero'}
     for group_name in group_set:
@@ -462,11 +460,11 @@ async def test_group_scope_filter(app):
 
 async def test_vertical_filter(app):
     user_name = 'lindsey'
-    add_user(app.db, name=user_name)
+    user = add_user(app.db, name=user_name)
     test_role = generate_test_role(user_name, ['read:users:name'])
-    roles.add_role(app.db, test_role)
+    roles.create_role(app.db, test_role)
-    roles.add_obj(app.db, objname=user_name, kind='users', rolename='test')
+    roles.grant_role(app.db, entity=user, rolename='test')
-    roles.remove_obj(app.db, objname=user_name, kind='users', rolename='user')
+    roles.strip_role(app.db, entity=user, rolename='user')
     app.db.commit()
 
     r = await api_request(app, 'users', headers=auth_header(app.db, user_name))
@@ -477,12 +475,13 @@ async def test_vertical_filter(app):
 
 async def test_stacked_vertical_filter(app):
     user_name = 'user'
+    user = add_user(app.db, name=user_name)
     test_role = generate_test_role(
         user_name, ['read:users:activity', 'read:users:servers']
     )
-    roles.add_role(app.db, test_role)
+    roles.create_role(app.db, test_role)
-    roles.add_obj(app.db, objname=user_name, kind='users', rolename='test')
+    roles.grant_role(app.db, entity=user, rolename='test')
-    roles.remove_obj(app.db, objname=user_name, kind='users', rolename='user')
+    roles.strip_role(app.db, entity=user, rolename='user')
     app.db.commit()
 
     r = await api_request(app, 'users', headers=auth_header(app.db, user_name))
@@ -494,13 +493,13 @@ async def test_stacked_vertical_filter(app):
 
 async def test_cross_filter(app):
     user_name = 'abed'
-    add_user(app.db, name=user_name)
+    user = add_user(app.db, name=user_name)
     test_role = generate_test_role(
         user_name, ['read:users:activity', 'read:users!user=abed']
     )
-    roles.add_role(app.db, test_role)
+    roles.create_role(app.db, test_role)
-    roles.add_obj(app.db, objname=user_name, kind='users', rolename='test')
+    roles.grant_role(app.db, entity=user, rolename='test')
-    roles.remove_obj(app.db, objname=user_name, kind='users', rolename='user')
+    roles.strip_role(app.db, entity=user, rolename='user')
     app.db.commit()
     new_users = {'britta', 'jeff', 'annie'}
     for new_user_name in new_users:

jupyterhub/tests/test_services.py

@@ -95,7 +95,7 @@ async def test_external_service(app):
 
         service = app._service_map[name]
         api_token = service.orm.api_tokens[0]
-        update_roles(app.db, api_token, 'tokens', roles=['token'])
+        update_roles(app.db, api_token, roles=['token'])
         url = public_url(app, service) + '/api/users'
         r = await async_requests.get(url, allow_redirects=False)
         r.raise_for_status()

jupyterhub/tests/utils.py

@@ -9,6 +9,7 @@ from certipy import Certipy
 from jupyterhub import metrics
 from jupyterhub import orm
 from jupyterhub.objects import Server
+from jupyterhub.roles import assign_default_roles
 from jupyterhub.roles import update_roles
 from jupyterhub.utils import url_path_join as ujoin
 
@@ -113,7 +114,10 @@ def add_user(db, app=None, **kwargs):
             setattr(orm_user, attr, value)
     db.commit()
     requested_roles = kwargs.get('roles')
-    update_roles(db, obj=orm_user, kind='users', roles=requested_roles)
+    if requested_roles:
+        update_roles(db, entity=orm_user, roles=requested_roles)
+    else:
+        assign_default_roles(db, entity=orm_user)
     if app:
         return app.users[orm_user.id]
     else: