From 70b11d9015f7f3b5b8bb31fd4b86c536e5fddf2e Mon Sep 17 00:00:00 2001
From: Manuel Ebert <maebert@uos.de>
Date: Fri, 25 Jan 2013 16:53:01 -0800
Subject: [PATCH] Unit-tests for flask-peewee

---
 .travis.yml                  |  2 +-
 setup.py                     |  1 +
 tests/functional_tests.py    |  7 ++++
 tests/test_app/peewee_app.py | 62 ++++++++++++++++++++++++++++++++++++
 4 files changed, 71 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_app/peewee_app.py

diff --git a/.travis.yml b/.travis.yml
index 42c3b47..a182d2f 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -7,7 +7,7 @@ python:
 install:
   - pip install . --quiet --use-mirrors
   - "if [[ $TRAVIS_PYTHON_VERSION != '2.7' ]]; then pip install importlib --quiet --use-mirrors; fi"
-  - pip install nose simplejson Flask-SQLAlchemy Flask-MongoEngine Flask-Mail py-bcrypt MySQL-python --quiet --use-mirrors
+  - pip install nose simplejson Flask-SQLAlchemy Flask-MongoEngine Flask-Peewee Flask-Mail py-bcrypt MySQL-python --quiet --use-mirrors
 
 before_script:
   - mysql -e 'create database flask_security_test;'
diff --git a/setup.py b/setup.py
index 4025e75..05e8986 100644
--- a/setup.py
+++ b/setup.py
@@ -47,6 +47,7 @@ setup(
         'nose',
         'Flask-SQLAlchemy',
         'Flask-MongoEngine',
+        'Flask-Peewee',
         'py-bcrypt',
         'simplejson'
     ],
diff --git a/tests/functional_tests.py b/tests/functional_tests.py
index c68625a..fd8d60a 100644
--- a/tests/functional_tests.py
+++ b/tests/functional_tests.py
@@ -224,6 +224,13 @@ class MongoEngineSecurityTests(DefaultSecurityTests):
         return create_app(auth_config, **kwargs)
 
 
+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)
+
+
 class DefaultDatastoreTests(SecurityTest):
 
     def test_add_role_to_user(self):
diff --git a/tests/test_app/peewee_app.py b/tests/test_app/peewee_app.py
new file mode 100644
index 0000000..ae404a2
--- /dev/null
+++ b/tests/test_app/peewee_app.py
@@ -0,0 +1,62 @@
+# -*- coding: utf-8 -*-
+
+import sys
+import os
+
+sys.path.pop(0)
+sys.path.insert(0, os.getcwd())
+
+from flask_peewee.db import Database
+from peewee import *
+from flask.ext.security import Security, UserMixin, RoleMixin, \
+     PeeweeUserDatastore
+
+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',
+    }
+    db = Database(app)
+
+    class Role(db.Model, RoleMixin):
+        name = TextField(unique=True)
+        description = TextField(null=True)
+
+    class User(db.Model, UserMixin):
+        email = TextField()
+        password = TextField()
+        last_login_at = DateTimeField(null=True)
+        current_login_at = DateTimeField(null=True)
+        last_login_ip = TextField(null=True)
+        current_login_ip = TextField(null=True)
+        login_count = IntegerField(null=True)
+        active = BooleanField(default=True)
+        confirmed_at = DateTimeField(null=True)
+
+    class UserRoles(db.Model):
+        """ Peewee does not have built-in many-to-many support, so we have to
+        create this mapping class to link users to roles."""
+        user = ForeignKeyField(User, related_name='roles')
+        role = ForeignKeyField(Role, related_name='users')
+        name = property(lambda self: self.role.name)
+        description = property(lambda self: self.role.description)
+
+    @app.before_first_request
+    def before_first_request():
+        for Model in (Role, User, UserRoles):
+            Model.drop_table(fail_silently=True)
+            Model.create_table(fail_silently=True)
+        populate_data(app.config.get('USER_COUNT', None))
+
+    app.security = Security(app, datastore=PeeweeUserDatastore(db, User, Role, UserRoles), **kwargs)
+
+    add_context_processors(app.security)
+
+    return app
+
+if __name__ == '__main__':
+    create_app({}).run()