diff --git a/test/test_wsgi_compliance.py b/test/test_wsgi_compliance.py index 54c356f..0df822a 100644 --- a/test/test_wsgi_compliance.py +++ b/test/test_wsgi_compliance.py @@ -85,3 +85,12 @@ def test_encoding_errors(): response, content = http.request( 'http://some_hopefully_nonexistant_domain/boom/baz', headers={'Accept': u'application/\u2603'}) + + +def test_post_status_headers(): + with InstalledApp(wsgi_app.post_status_headers_app, host=HOST) as app: + http = httplib2.Http() + resp, content = http.request( + 'http://some_hopefully_nonexistant_domain/', 'GET') + assert app.success() + assert resp.get('content-type') == 'text/plain' diff --git a/test/wsgi_app.py b/test/wsgi_app.py index aecf1ed..f0c4225 100644 --- a/test/wsgi_app.py +++ b/test/wsgi_app.py @@ -23,5 +23,12 @@ def more_interesting_app(environ, start_response): return [pformat(environ).encode('utf-8')] +def post_status_headers_app(environ, start_response): + headers = [] + start_response('200 OK', headers) + headers.append(('Content-type', 'text/plain')) + return [b'WSGI intercept successful!\n'] + + def raises_app(environ, start_response): raise TypeError("bah") diff --git a/wsgi_intercept/__init__.py b/wsgi_intercept/__init__.py index 3772842..369ca0e 100644 --- a/wsgi_intercept/__init__.py +++ b/wsgi_intercept/__init__.py @@ -375,25 +375,15 @@ class wsgi_fake_socket: # dynamically construct the start_response function for no good reason. + self.headers = [] + def start_response(status, headers, exc_info=None): # construct the HTTP request. self.output.write(b"HTTP/1.0 " + status.encode('utf-8') + b"\n") - - for k, v in headers: - try: - k = k.encode('utf-8') - except AttributeError: - pass - try: - v = v.encode('utf-8') - except AttributeError: - pass - self.output.write(k + b': ' + v + b"\n") - self.output.write(b'\n') - - def write_fn(s): - self.write_results.append(s) - return write_fn + # Keep the reference of the headers list to write them only + # when the whole application have been processed + self.headers = headers + return self.write_results.append # construct the wsgi.input file from everything that's been # written to this "socket". @@ -411,6 +401,20 @@ class wsgi_fake_socket: raise WSGIAppError(error, sys.exc_info()) self.result = iter(app_result) + # send the headers + + for k, v in self.headers: + try: + k = k.encode('utf-8') + except AttributeError: + pass + try: + v = v.encode('utf-8') + except AttributeError: + pass + self.output.write(k + b': ' + v + b"\n") + self.output.write(b'\n') + ### # read all of the results. the trick here is to get the *first*