From 38874433c7fe2cc570b40a7e99950b2d25b4315e Mon Sep 17 00:00:00 2001
From: Matt Wright <matt@nobien.net>
Date: Thu, 4 Apr 2013 18:09:55 -0400
Subject: [PATCH] Add tests for Peewee support

---
 flask_security/datastore.py  | 14 +++++---------
 tests/functional_tests.py    |  8 ++++----
 tests/test_app/peewee_app.py |  3 ++-
 3 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/flask_security/datastore.py b/flask_security/datastore.py
index a08d5ae..e3e600a 100644
--- a/flask_security/datastore.py
+++ b/flask_security/datastore.py
@@ -95,12 +95,11 @@ class UserDatastore(object):
         :param user: The user to manipulate
         :param role: The role to add to the user
         """
-        rv = False
         user, role = self._prepare_role_modify_args(user, role)
         if role not in user.roles:
-            rv = True
             user.roles.append(role)
-        return rv
+            return True
+        return False
 
     def remove_role_from_user(self, user, role):
         """Removes a role from a user
@@ -220,10 +219,7 @@ class PeeweeUserDatastore(PeeweeDatastore, UserDatastore):
             return None
 
     def find_role(self, role):
-        try:
-            return self.role_model.filter(name=role).get()
-        except self.role_model.DoesNotExist:
-            return None
+        return self.role_model.filter(name=role).get()
 
     def create_user(self, **kwargs):
         """Creates and returns a new user from the given parameters."""
@@ -242,11 +238,11 @@ class PeeweeUserDatastore(PeeweeDatastore, UserDatastore):
         """
         user, role = self._prepare_role_modify_args(user, role)
         result = self.UserRole.select() \
-            .where(self.UserRole.user == user, self.UserRole.role == role)
+            .where(self.UserRole.user == user.id, self.UserRole.role == role.id)
         if result.count():
             return False
         else:
-            self.UserRole.create(user=user, role=role)
+            self.UserRole.create(user=user.id, role=role.id)
             return True
 
     def remove_role_from_user(self, user, role):
diff --git a/tests/functional_tests.py b/tests/functional_tests.py
index 20edd2c..c808d6f 100644
--- a/tests/functional_tests.py
+++ b/tests/functional_tests.py
@@ -231,11 +231,11 @@ class MongoEngineSecurityTests(DefaultSecurityTests):
         return create_app(auth_config, **kwargs)
 
 
-# class PeeweeSecurityTests(DefaultSecurityTests):
+class PeeweeSecurityTests(DefaultSecurityTests):
 
-#     def _create_app(self, auth_config, **kwargs):
-#         from tests.test_app.peewee_app import create_app
-#         return create_app(auth_config, **kwargs)
+    def _create_app(self, auth_config, **kwargs):
+        from tests.test_app.peewee_app import create_app
+        return create_app(auth_config, **kwargs)
 
 
 class DefaultDatastoreTests(SecurityTest):
diff --git a/tests/test_app/peewee_app.py b/tests/test_app/peewee_app.py
index ae404a2..954a977 100644
--- a/tests/test_app/peewee_app.py
+++ b/tests/test_app/peewee_app.py
@@ -14,11 +14,12 @@ from flask.ext.security import Security, UserMixin, RoleMixin, \
 from tests.test_app import create_app as create_base_app, populate_data, \
      add_context_processors
 
+
 def create_app(config, **kwargs):
     app = create_base_app(config)
     app.config['DATABASE'] = {
         'name': 'example2.db',
-        'engine': 'peewee.SqliteDatabase',
+        'engine': 'peewee.SqliteDatabase'
     }
     db = Database(app)