From 40b1ec91f084080813e8b17d4fc20c308973bd6e Mon Sep 17 00:00:00 2001
From: Channel Cat <channelcat@gmail.com>
Date: Sun, 16 Oct 2016 01:48:51 -0700
Subject: [PATCH] Added blueprint exception and middleware support, and tests

---
 sanic/blueprints.py      | 57 +++++++++++++++++++++++++++++++---------
 tests/test_blueprints.py | 57 +++++++++++++++++++++++++++++++++++++---
 2 files changed, 98 insertions(+), 16 deletions(-)

diff --git a/sanic/blueprints.py b/sanic/blueprints.py
index 585863bb..506c2e90 100644
--- a/sanic/blueprints.py
+++ b/sanic/blueprints.py
@@ -1,6 +1,4 @@
-
-
-class BlueprintSetup():
+class BlueprintSetup:
     """
     """
 
@@ -17,26 +15,41 @@ class BlueprintSetup():
         #: blueprint.
         self.url_prefix = url_prefix
 
-    def add_url_rule(self, uri, methods=None, handler=None, **options):
-        """A helper method to register a handler to the application url routes.
-
+    def add_route(self, handler, uri, methods):
+        """
+        A helper method to register a handler to the application url routes.
         """
         if self.url_prefix:
             uri = self.url_prefix + uri
 
         self.app.router.add(uri, methods, handler)
 
+    def add_exception(self, handler, *args, **kwargs):
+        """
+        Registers exceptions to sanic
+        """
+        self.app.exception(*args, **kwargs)(handler)
 
-class Blueprint():
+    def add_middleware(self, middleware, *args, **kwargs):
+        """
+        Registers middleware to sanic
+        """
+        if args or kwargs:
+            self.app.middleware(*args, **kwargs)(middleware)
+        else:
+            self.app.middleware(middleware)
+
+
+class Blueprint:
     def __init__(self, name, url_prefix=None):
         self.name = name
         self.url_prefix = url_prefix
         self.deferred_functions = []
 
     def record(self, func):
-        """Registers a callback function that is invoked when the blueprint is
+        """
+        Registers a callback function that is invoked when the blueprint is
         registered on the application.
-
         """
         self.deferred_functions.append(func)
 
@@ -57,12 +70,30 @@ class Blueprint():
         """
         """
         def decorator(handler):
-            self.add_url_rule(uri=uri, methods=methods, handler=handler)
+            self.record(lambda s: s.add_route(handler, uri, methods))
             return handler
         return decorator
 
-    def add_url_rule(self, uri, methods=None, handler=None):
+    def middleware(self, *args, **kwargs):
         """
         """
-        self.record(lambda s:
-            s.add_url_rule(uri, methods, handler))
+
+        def register_middleware(middleware):
+            self.record(lambda s: s.add_middleware(middleware, *args, **kwargs))
+            return middleware
+
+        # Detect which way this was called, @middleware or @middleware('AT')
+        if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
+            args = []
+            return register_middleware(args[0])
+        else:
+            return register_middleware
+
+    def exception(self, *args, **kwargs):
+        """
+        """
+        def decorator(handler):
+            self.record(lambda s: s.add_exception(handler, *args, **kwargs))
+            return handler
+        return decorator
+
diff --git a/tests/test_blueprints.py b/tests/test_blueprints.py
index a88f2d19..8068160f 100644
--- a/tests/test_blueprints.py
+++ b/tests/test_blueprints.py
@@ -1,9 +1,8 @@
-from json import loads as json_loads, dumps as json_dumps
 from sanic import Sanic
-from sanic import Blueprint
+from sanic.blueprints import Blueprint
 from sanic.response import json, text
 from sanic.utils import sanic_endpoint_test
-from sanic.exceptions import SanicException
+from sanic.exceptions import NotFound, ServerError, InvalidUsage
 
 
 # ------------------------------------------------------------ #
@@ -58,3 +57,55 @@ def test_several_bp_with_url_prefix():
     request, response = sanic_endpoint_test(app, uri='/test2/')
     assert response.text == 'Hello2'
 
+
+def test_bp_middleware():
+    app = Sanic('test_middleware')
+    blueprint = Blueprint('test_middleware')
+
+    @blueprint.middleware('response')
+    async def process_response(request, response):
+        return text('OK')
+
+    @app.route('/')
+    async def handler(request):
+        return text('FAIL')
+
+    app.register_blueprint(blueprint)
+
+    request, response = sanic_endpoint_test(app)
+
+    assert response.status == 200
+    assert response.text == 'OK'
+
+def test_bp_exception_handler():
+    app = Sanic('test_middleware')
+    blueprint = Blueprint('test_middleware')
+
+    @blueprint.route('/1')
+    def handler_1(request):
+        raise InvalidUsage("OK")
+
+    @blueprint.route('/2')
+    def handler_2(request):
+        raise ServerError("OK")
+
+    @blueprint.route('/3')
+    def handler_3(request):
+        raise NotFound("OK")
+
+    @blueprint.exception(NotFound, ServerError)
+    def handler_exception(request, exception):
+        return text("OK")
+
+    app.register_blueprint(blueprint)
+
+    request, response = sanic_endpoint_test(app, uri='/1')
+    assert response.status == 400
+
+
+    request, response = sanic_endpoint_test(app, uri='/2')
+    assert response.status == 200
+    assert response.text == 'OK'
+
+    request, response = sanic_endpoint_test(app, uri='/3')
+    assert response.status == 200
\ No newline at end of file