diff options
| -rw-r--r-- | raven/middleware.py | 47 | ||||
| -rw-r--r-- | tests/middleware/tests.py | 34 |
2 files changed, 62 insertions, 19 deletions
diff --git a/raven/middleware.py b/raven/middleware.py index 588bd6f..4e81a45 100644 --- a/raven/middleware.py +++ b/raven/middleware.py @@ -9,6 +9,37 @@ from __future__ import absolute_import from raven.utils.wsgi import ( get_current_url, get_headers, get_environ) +from raven.utils.six import Iterator, next + + +class ClosingIterator(Iterator): + """ + An iterator that is implements a ``close`` method as-per + WSGI recommendation. + """ + def __init__(self, sentry, iterable, environ): + self.sentry = sentry + self.environ = environ + self.iterable = iter(iterable) + + def __iter__(self): + return self + + def __next__(self): + try: + return next(self.iterable) + except StopIteration: + # propagate up the normal StopIteration + raise + except Exception: + # but capture any other exception, then re-raise + self.sentry.handle_exception(self.environ) + raise + + def close(self): + self.sentry.client.context.clear() + if hasattr(self.iterable, 'close') and callable(self.iterable.close): + self.iterable.close() class Sentry(object): @@ -37,21 +68,7 @@ class Sentry(object): self.handle_exception(environ) raise - try: - for event in iterable: - yield event - except Exception: - self.handle_exception(environ) - raise - finally: - # wsgi spec requires iterable to call close if it exists - # see http://blog.dscpl.com.au/2012/10/obligations-for-calling-close-on.html - if iterable and hasattr(iterable, 'close') and callable(iterable.close): - try: - iterable.close() - except Exception: - self.handle_exception(environ) - self.client.context.clear() + return ClosingIterator(self, iterable, environ) def get_http_context(self, environ): return { diff --git a/tests/middleware/tests.py b/tests/middleware/tests.py index 23e94ca..2a0beb7 100644 --- a/tests/middleware/tests.py +++ b/tests/middleware/tests.py @@ -7,6 +7,7 @@ from raven.utils.testutils import TestCase from raven.base import Client from raven.middleware import Sentry +from raven.utils.six import Iterator, next class TempStoreClient(Client): @@ -21,17 +22,35 @@ class TempStoreClient(Client): self.events.append(kwargs) -class ErroringIterable(object): +class ErroringIterable(Iterator): def __init__(self): self.closed = False def __iter__(self): + return self + + def __next__(self): raise ValueError('hello world') def close(self): self.closed = True +class SimpleIteratable(Iterator): + def __init__(self): + self.closed = False + self._iter = iter(['a']) + + def __iter__(self): + return self + + def __next__(self): + return next(self._iter) + + def close(self): + self.closed = True + + class ExampleApp(object): def __init__(self, iterable): self.iterable = iterable @@ -59,9 +78,6 @@ class MiddlewareTestCase(TestCase): with self.assertRaises(ValueError): response = list(response) - # TODO: this should be a separate test - self.assertTrue(iterable.closed, True) - self.assertEquals(len(self.client.events), 1) event = self.client.events.pop(0) @@ -86,3 +102,13 @@ class MiddlewareTestCase(TestCase): self.assertEquals(env['SERVER_NAME'], 'localhost') self.assertTrue('SERVER_PORT' in env, env.keys()) self.assertEquals(env['SERVER_PORT'], '80') + + def test_close(self): + iterable = SimpleIteratable() + app = ExampleApp(iterable) + middleware = Sentry(app, client=self.client) + + response = middleware(self.request.environ, lambda *args: None) + list(response) # exhaust iterator + response.close() + self.assertTrue(iterable.closed, True) |