Role helpers and tests

src/database.py

@@ -1,7 +1,7 @@
 from pathlib import Path
 from sqlalchemy import *
 from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import sessionmaker, relationship
 from sqlalchemy.orm.exc import NoResultFound, MultipleResultsFound
 import os
 import os.path
@@ -279,13 +279,36 @@ class User(Base, BaseModel):
 
     UniqueConstraint(cso_user_id, cse_user_id, cmse_user_id)
 
+    user_roles = relationship('UserRole', back_populates='user')
+
+    _role_names = None
+
+    def role_names(self, reload=False):
+        if self._role_names is None or reload is True:
+            self._role_names = [x.role.role_name for x in self.user_roles]
+
+        return self._role_names
+
+    def has_role(self, role_name, reload=False):
+        return role_name in self.role_names(reload=reload)
+
 
 class Role(Base, BaseModel):
     __tablename__ = 'roles'
 
     id = Column(Integer, primary_key=True)
     role_name = Column(String(100))
 
+    user_roles = relationship('UserRole', back_populates='role')
+
+    _users = None
+
+    def users(self, reload=False):
+        if self._users is None or reload is True:
+            self._users = [x.user for x in self.user_roles]
+
+        return self._users
+
 
 class UserRole(Base, BaseModel):
     __tablename__ = 'users_roles'
@@ -294,3 +317,6 @@ class UserRole(Base, BaseModel):
     role_id = Column(Integer, ForeignKey(Role.id), nullable=False)
 
     PrimaryKeyConstraint(user_id, role_id)
+
+    user = relationship(User, back_populates='user_roles')
+    role = relationship(Role, back_populates='user_roles')

test/test_database.py

@@ -1,4 +1,5 @@
 from datetime import datetime
+from sqlalchemy import or_
 from src.database import *
 
 
@@ -21,16 +22,36 @@ def setup_module(module):
                      site_url='ng-test-ds.stackexchange.com'),
         SchemaMigration(migration_file='ng_test_d20170429021826.py', run_status=False, run_at=None),
         SchemaMigration(migration_file='ng_test_d20170429021856.py', run_status=True,
-                        run_at=datetime(2017, 4, 29, 2, 19, 26))
+                        run_at=datetime(2017, 4, 29, 2, 19, 26)),
+        User(cso_user_id=-999, cse_user_id=-999, cmse_user_id=-999),
+        Role(role_name='ng-test-admin')
     ]
     SESSION.add_all(records)
     SESSION.commit()
 
+    # Annoyingly, we have to do this twice so that we have access to the id attribute for link models.
+    test_user = SESSION.query(User).filter(User.cso_user_id == -999).first()
+    test_role = SESSION.query(Role).filter(Role.role_name == 'ng-test-admin').first()
+    link_records = [
+        UserRole(user_id=test_user.id, role_id=test_role.id)
+    ]
+    SESSION.add_all(link_records)
+    SESSION.commit()
+
 
 def teardown_module(module):
     SESSION.query(Notification).filter(Notification.site_url.like('ng-test%')).delete(synchronize_session=False)
     SESSION.query(SchemaMigration).filter(SchemaMigration.migration_file.like('ng_test%'))\
            .delete(synchronize_session=False)
+
+    user_ids = [x.id for x in SESSION.query(User).filter(User.cso_user_id == -999).all()]
+    role_ids = [x.id for x in SESSION.query(Role).filter(Role.role_name.like('ng-test%')).all()]
+    SESSION.query(UserRole).filter(or_(UserRole.user_id.in_(user_ids), UserRole.role_id.in_(role_ids)))\
+           .delete(synchronize_session=False)
+
+    SESSION.query(User).filter(User.id.in_(user_ids)).delete(synchronize_session=False)
+    SESSION.query(Role).filter(Role.id.in_(role_ids)).delete(synchronize_session=False)
+
     SESSION.commit()
 
 
@@ -95,3 +116,19 @@ def test_schema_migration_pending():
 def test_schema_migration_is_run():
     assert SchemaMigration.is_run('ng_test_d20170429021856.py') is True
     assert SchemaMigration.is_run('ng_test_d20170429021826.py') is False
+
+
+def test_user_role_names():
+    user = SESSION.query(User).filter(User.cso_user_id == -999).first()
+    assert user.role_names() == ['ng-test-admin']
+
+
+def test_user_has_role():
+    user = SESSION.query(User).filter(User.cso_user_id == -999).first()
+    assert user.has_role('ng-test-admin')
+
+
+def test_role_users():
+    role = SESSION.query(Role).filter(Role.role_name == 'ng-test-admin').first()
+    user = SESSION.query(User).filter(User.cso_user_id == -999).first()
+    assert role.users() == [user]