New User role management

hitier
1 parent a5a365e8
Showing 2 changed files with 65 additions and 1 deletions Show diff stats
app/auth/models.py
tests/backend_tests.py
 from flask_login import UserMixin
 from app.models import db
  
+#
+# Roles
+#
+
+ADMIN = 40
+SERVICE = 30
+AGENT = 10
+PUBLIC = 0
+
+_roleToName = {
+    ADMIN: 'ADMIN',
+    SERVICE: 'SERVICE',
+    AGENT: 'AGENT',
+    PUBLIC: 'PUBLIC',
+}
+_nameToRole = {
+    'ADMIN': ADMIN,
+    'SERVICE': SERVICE,
+    'AGENT': AGENT,
+    'PUBLIC': PUBLIC,
+}
+
+
+def _checkRole(role):
+    if isinstance(role, int):
+        rv = role
+    elif str(role) == role:
+        role = role.upper()
+        if role not in _nameToRole:
+            raise ValueError("Unknown role: %r" % role)
+        rv = _nameToRole[role]
+    else:
+        raise TypeError("Role not an integer or a valid string: %r" % role)
+    return rv
+
  
 class User(UserMixin, db.Model):
     id = db.Column(db.Integer, primary_key=True)  # primary keys are required by SQLAlchemy
@@ -8,7 +43,14 @@ class User(UserMixin, db.Model):
     name = db.Column(db.String(100))
     login = db.Column(db.String(100), unique=True)
     password = db.Column(db.String(100))
+    role = db.Column(db.Integer)
  
     def __repr__(self):
         return "i: {}, n: {}, e: {}, l: {}".format(self.id, self.name, self.email, self.login)
  
+    def setRole(self, role):
+        self.role = _checkRole(role)
+
+    def hasRole(self, role):
+        role = _checkRole(role)
+        return role == self.role
 import unittest
 from pdc_config import TestConfig
-from app import create_app, db_mgr
+from app import create_app, db_mgr, db
+from app.auth.models import User
  
  
 class BaseTestCase(unittest.TestCase):
@@ -9,6 +10,10 @@ class BaseTestCase(unittest.TestCase):
         self.app = create_app(TestConfig)
         self.app_context = self.app.app_context()
         self.app_context.push()
+        db.create_all()
+        admin = User(email='admin@nowhere.org', name='admin', login='admin', password='admin')
+        db.session.add(admin)
+        db.session.commit()
  
     def tearDown(self):
         self.app_context.pop()
@@ -27,3 +32,20 @@ class DbMgrTestCase(BaseTestCase):
     def test_charges_by_agent(self):
         all_charges = db_mgr.charges_by_agent(355)
         self.assertEqual(6, len(all_charges))
+
+
+class AuthModelTestCase(BaseTestCase):
+
+    def test_setrole(self):
+        admin = User.query.filter(User.name == 'admin').one_or_none()
+        admin.setRole("ADMIN")
+        db.session.commit()
+        admin = User.query.filter(User.name == 'admin').one_or_none()
+        self.assertTrue(admin is not None)
+        self.assertTrue(admin.hasRole("ADMIN"))
+        self.assertFalse(admin.hasRole("SERVICE"))
+
+    def test_setrole_valueerror(self):
+        admin = User(email='me@nowhere.org', name='me', login='me', password='me')
+        with self.assertRaises(ValueError) as ve:
+            admin.setRole("NOSUCHROLE")