From b0b09aea4913efcde116ef8cc1da03790ae17912 Mon Sep 17 00:00:00 2001
From: Matt Wright <mdw1980@gmail.com>
Date: Thu, 23 Aug 2012 13:01:11 -0400
Subject: [PATCH] Add ability to define a send_mail_task which could be used to
 send mails instead of the default flask-mail plugin. Could also be used to
 send mail asynchronously. Make flask-mail required as well.

---
 flask_security/core.py    | 13 ++++++++-----
 flask_security/utils.py   |  9 +++++++--
 setup.py                  |  1 +
 tests/functional_tests.py | 19 +++++++++++++++++++
 4 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/flask_security/core.py b/flask_security/core.py
index 2d4be43..eb7d5c5 100644
--- a/flask_security/core.py
+++ b/flask_security/core.py
@@ -210,6 +210,7 @@ class _SecurityState(object):
     def __init__(self, **kwargs):
         for key, value in kwargs.items():
             setattr(self, key.lower(), value)
+        self._send_mail_task = None
 
     def _add_ctx_processor(self, endpoint, fn):
         c = self.context_processors
@@ -256,6 +257,9 @@ class _SecurityState(object):
     def mail_context_processor(self, fn):
         self._add_ctx_processor('mail', fn)
 
+    def send_mail_task(self, fn):
+        self._send_mail_task = fn
+
 
 class Security(object):
     """The :class:`Security` class initializes the Flask-Security extension.
@@ -270,7 +274,7 @@ class Security(object):
         if app is not None and datastore is not None:
             self._state = self.init_app(app, datastore, **kwargs)
 
-    def init_app(self, app, datastore=None, register_blueprint=True):
+    def init_app(self, app, datastore=None, register_blueprint=True, **kwargs):
         """Initializes the Flask-Security extension for the specified
         application and datastore implentation.
 
@@ -295,7 +299,7 @@ class Security(object):
                                         template_folder='templates')
             app.register_blueprint(bp)
 
-        state = self._get_state(app, datastore)
+        state = self._get_state(app, datastore, **kwargs)
 
         app.extensions['security'] = state
 
@@ -304,12 +308,10 @@ class Security(object):
 
         return state
 
-    def _get_state(self, app, datastore):
+    def _get_state(self, app, datastore, **kwargs):
         assert app is not None
         assert datastore is not None
 
-        kwargs = {}
-
         for key, value in get_config(app).items():
             kwargs[key.lower()] = value
 
@@ -329,6 +331,7 @@ class Security(object):
             _get_reset_serializer(app) if kwargs['recoverable'] else None)
         kwargs['confirm_serializer'] = (
             _get_confirm_serializer(app) if kwargs['confirmable'] else None)
+
         return _SecurityState(**kwargs)
 
     def __getattr__(self, name):
diff --git a/flask_security/utils.py b/flask_security/utils.py
index a963877..3030290 100644
--- a/flask_security/utils.py
+++ b/flask_security/utils.py
@@ -20,6 +20,7 @@ from flask import url_for, flash, current_app, request, session, redirect, \
      render_template
 from flask.ext.login import login_user as _login_user, \
      logout_user as _logout_user
+from flask.ext.mail import Message
 from flask.ext.principal import Identity, AnonymousIdentity, identity_changed
 from werkzeug.local import LocalProxy
 
@@ -238,9 +239,7 @@ def send_mail(subject, recipient, template, **context):
     :param template: The name of the email template
     :param context: The context to render the template with
     """
-    from flask.ext.mail import Message
 
-    mail = current_app.extensions.get('mail')
     context.setdefault('security', _security)
     context.update(_security._run_ctx_processor('mail'))
 
@@ -251,6 +250,12 @@ def send_mail(subject, recipient, template, **context):
     ctx = ('security/email', template)
     msg.body = render_template('%s/%s.txt' % ctx, **context)
     msg.html = render_template('%s/%s.html' % ctx, **context)
+
+    if _security._send_mail_task:
+        _security._send_mail_task(msg)
+        return
+
+    mail = current_app.extensions.get('mail')
     mail.send(msg)
 
 
diff --git a/setup.py b/setup.py
index c24e762..45fa170 100644
--- a/setup.py
+++ b/setup.py
@@ -36,6 +36,7 @@ setup(
     install_requires=[
         'Flask>=0.9',
         'Flask-Login>=0.1.3',
+        'Flask-Mail>=0.6.1',
         'Flask-Principal>=0.3',
         'Flask-WTF>=0.5.4',
         'itsdangerous>=0.15',
diff --git a/tests/functional_tests.py b/tests/functional_tests.py
index 9752f17..48d6e5c 100644
--- a/tests/functional_tests.py
+++ b/tests/functional_tests.py
@@ -590,3 +590,22 @@ class MongoEngineDatastoreTests(DefaultDatastoreTests):
     def _create_app(self, auth_config):
         from tests.test_app.mongoengine import create_app
         return create_app(auth_config)
+
+
+class AsyncMailTaskTests(SecurityTest):
+
+    AUTH_CONFIG = {
+        'SECURITY_RECOVERABLE': True,
+    }
+
+    def setUp(self):
+        super(AsyncMailTaskTests, self).setUp()
+        self.mail_sent = False
+
+    def test_send_email_task_is_called(self):
+        @self.app.security.send_mail_task
+        def send_email(msg):
+            self.mail_sent = True
+
+        self.client.post('/reset', data=dict(email='joe@lp.com'))
+        self.assertTrue(self.mail_sent)