diff --git a/app/auth/models.py b/app/auth/models.py index 42b0066..e5194e0 100644 --- a/app/auth/models.py +++ b/app/auth/models.py @@ -1,6 +1,41 @@ 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 diff --git a/tests/backend_tests.py b/tests/backend_tests.py index cb9f245..cc78bf8 100644 --- a/tests/backend_tests.py +++ b/tests/backend_tests.py @@ -1,6 +1,7 @@ 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") -- libgit2 0.21.2