diff options
| -rw-r--r-- | paste/session.py | 19 | ||||
| -rw-r--r-- | tests/test_session.py | 52 |
2 files changed, 67 insertions, 4 deletions
diff --git a/paste/session.py b/paste/session.py index b10601f..c2feac6 100644 --- a/paste/session.py +++ b/paste/session.py @@ -42,18 +42,29 @@ class SessionMiddleware(object): def __call__(self, environ, start_response): session_factory = SessionFactory(environ, **self.factory_kw) environ['paste.session.factory'] = session_factory + remember_headers = [] def session_start_response(status, headers, exc_info=None): if not session_factory.created: + remember_headers[:] = [status, headers] return start_response(status, headers) headers.append(session_factory.set_cookie_header()) return start_response(status, headers, exc_info) app_iter = self.application(environ, session_start_response) - if session_factory.used: - return wsgilib.add_close(app_iter, session_factory.close) - else: - return app_iter + def start(): + if session_factory.created and remember_headers: + # Tricky bastard used the session after start_response + status, headers = remember_headers + headers.append(session_factory.set_cookie_header()) + exc = ValueError( + "You cannot get the session after content from the " + "app_iter has been returned") + start_response(status, headers, (exc.__class__, exc, None)) + def close(): + if session_factory.used: + session_factory.close() + return wsgilib.add_start_close(app_iter, start, close) class SessionFactory(object): diff --git a/tests/test_session.py b/tests/test_session.py new file mode 100644 index 0000000..621d284 --- /dev/null +++ b/tests/test_session.py @@ -0,0 +1,52 @@ +from paste.session import SessionMiddleware +from paste.fixture import TestApp + +info = [] + +def wsgi_app(environ, start_response): + pi = environ.get('PATH_INFO', '') + if pi in ('/get1', '/get2'): + if pi == '/get1': + sess = environ['paste.session.factory']() + start_response('200 OK', [('content-type', 'text/plain')]) + if pi == '/get2': + sess = environ['paste.session.factory']() + if 'info' in sess: + return [str(sess['info'])] + else: + return ['no-info'] + if pi in ('/put1', '/put2'): + if pi == '/put1': + sess = environ['paste.session.factory']() + sess['info'] = info[0] + start_response('200 OK', [('content-type', 'text/plain')]) + if pi == '/put2': + sess = environ['paste.session.factory']() + sess['info'] = info[0] + return ['foo'] + +wsgi_app = SessionMiddleware(wsgi_app) + +def test_app1(): + app = TestApp(wsgi_app) + res = app.get('/get1') + assert res.body == 'no-info' + res = app.get('/get2') + assert res.body == 'no-info' + info[:] = ['test'] + res = app.get('/put1') + res = app.get('/get1') + assert res.body == 'test' + res = app.get('/get2') + assert res.body == 'test' + +def test_app2(): + app = TestApp(wsgi_app) + info[:] = ['fluff'] + res = app.get('/put2') + res = app.get('/get1') + assert res.body == 'fluff' + res = app.get('/get2') + assert res.body == 'fluff' + + |