Tests and bugfixes for cgitb_catcher

2 files changed, 106 insertions, 20 deletions

diff --git a/paste/cgitb_catcher.py b/paste/cgitb_catcher.py
index a245299..43df21c 100644
--- a/paste/cgitb_catcher.py
+++ b/paste/cgitb_catcher.py
@@ -22,7 +22,7 @@ class CgitbMiddleware(object):
                  logdir=None,
                  context=5,
                  format="html"):
-        self.app = application
+        self.app = app
         if display is NoDefault:
             display = global_conf.get('debug')
         self.display = converters.asbool(display)
@@ -33,34 +33,45 @@ class CgitbMiddleware(object):
     def __call__(self, environ, start_response):
         try:
             app_iter = self.app(environ, start_response)
-            return catching_iter(app_iter)
+            return self.catching_iter(app_iter, environ)
         except:
             exc_info = sys.exc_info()
             start_response('500 Internal Server Error',
                            [('content-type', 'text/html')],
                            exc_info)
-            dummy_file = StringIO()
-            hook = cgitb.Hook(file=dummy_file,
-                              display=self.display,
-                              logdir=self.logdir,
-                              context=self.context,
-                              format=self.format)
-            hook(*exc_info)
-            return [dummy_file.getvalue()]
+            response = self.exception_handler(exc_info, environ)
+            return [response]
 
-    def catching_iter(self, app_iter):
+    def catching_iter(self, app_iter, environ):
         if not app_iter:
             raise StopIteration
+        error_on_close = False
         try:
             for v in app_iter:
                 yield v
+            if hasattr(app_iter, 'close'):
+                error_on_close = True
+                app_iter.close()
         except:
-            exc_info = sys.exc_info()
-            dummy_file = StringIO()
-            hook = cgitb.Hook(file=dummy_file,
-                              display=self.display,
-                              logdir=self.logdir,
-                              context=self.context,
-                              format=self.format)
-            hook(*exc_info)
-            yield dummy_file.getvalue()
+            response = self.exception_handler(sys.exc_info(), environ)
+            if not error_on_close and hasattr(app_iter, 'close'):
+                try:
+                    app_iter.close()
+                except:
+                    close_response = self.exception_handler(
+                        sys.exc_info(), environ)
+                    response += (
+                        '<hr noshade>Error in .close():<br>%s'
+                        % close_response)
+            yield response
+
+    def exception_handler(self, exc_info, environ):
+        dummy_file = StringIO()
+        hook = cgitb.Hook(file=dummy_file,
+                          display=self.display,
+                          logdir=self.logdir,
+                          context=self.context,
+                          format=self.format)
+        hook(*exc_info)
+        return dummy_file.getvalue()
+        
diff --git a/tests/test_cgitb_catcher.py b/tests/test_cgitb_catcher.py
new file mode 100644
index 0000000..2b80d28
--- /dev/null
+++ b/tests/test_cgitb_catcher.py
@@ -0,0 +1,75 @@
+from paste.fixture import *
+from paste.cgitb_catcher import CgitbMiddleware
+from paste import lint
+from test_exceptions.test_error_middleware import clear_middleware
+
+def do_request(app, expect_status=500):
+    app = lint.middleware(app)
+    app = CgitbMiddleware(app, {}, display=True)
+    app = clear_middleware(app)
+    testapp = TestApp(app)
+    res = testapp.get('', status=expect_status,
+                      expect_errors=True)
+    return res
+    
+
+############################################################
+## Applications that raise exceptions
+############################################################
+
+def bad_app():
+    "No argument list!"
+    return None
+
+def start_response_app(environ, start_response):
+    "raise error before start_response"
+    raise ValueError("hi")
+
+def after_start_response_app(environ, start_response):
+    start_response("200 OK", [('Content-type', 'text/plain')])
+    raise ValueError('error2')
+
+def iter_app(environ, start_response):
+    start_response("200 OK", [('Content-type', 'text/plain')])
+    return yielder(['this', ' is ', ' a', None])
+
+def yielder(args):
+    for arg in args:
+        if arg is None:
+            raise ValueError("None raises error")
+        yield arg
+
+############################################################
+## Tests
+############################################################
+
+def test_makes_exception():
+    res = do_request(bad_app)
+    print res
+    assert 'bad_app() takes no arguments (2 given' in res
+    assert 'iterator = application(environ, start_response_wrapper)' in res
+    assert 'lint.py' in res
+    assert 'cgitb_catcher.py' in res
+
+def test_start_res():
+    res = do_request(start_response_app)
+    print res
+    assert 'ValueError: hi' in res
+    assert 'test_cgitb_catcher.py' in res
+    assert 'line 26, in start_response_app' in res
+
+def test_after_start():
+    res = do_request(after_start_response_app, 200)
+    print res
+    assert 'ValueError: error2' in res
+    assert 'line 30' in res
+
+def test_iter_app():
+    res = do_request(iter_app, 200)
+    print res
+    assert 'None raises error' in res
+    assert 'yielder' in res
+    
+                      
+
+