Modify after request hook

kobin/app.py

@@ -66,7 +66,9 @@ def _handle(self, environ):
             response = callback(**kwargs) if kwargs else callback()
 
             if self.after_request_callback:
-                response = self.after_request_callback(response)
+                wrapped_response = self.after_request_callback(response)
+                if wrapped_response:
+                    response = wrapped_response
         except HTTPError as e:
             response = e
         except BaseException as e:

tests/test_apps.py

@@ -117,6 +117,26 @@ def test_after_request(self):
         self.assertIn(('Foo', 'Bar'), response.headerlist)
 
 
+class KobinAfterHookTests(TestCase):
+    def setUp(self):
+        self.app = Kobin()
+        self.dummy_start_response = lambda x, y: None
+        self.before_counter = 0
+
+        @self.app.route('/')
+        def dummy_func():
+            return Response('hello')
+
+        @self.app.after_request
+        def after_do_not_return_response(response):
+            pass
+
+    def test_after_request(self):
+        test_env = {'REQUEST_METHOD': 'GET', 'PATH_INFO': '/'}
+        response = self.app._handle(test_env)
+        self.assertEqual('200 OK', response.status)
+
+
 class ConfigTests(TestCase):
     def setUp(self):
         self.root_path = os.path.dirname(os.path.abspath(__file__))