From 585f919c4bf8cc10b68703856ebab7a3b64ba1a3 Mon Sep 17 00:00:00 2001
From: Mehdi Abaakouk <sileht@sileht.net>
Date: Fri, 31 Jul 2015 14:20:29 +0200
Subject: [PATCH] process headers only when application is processed

An application can modify the headers after
start_response have been processed.

So we have to store the references of the headers list.

And write them to the client only when the application
have been processed but before we iter the body.

wsgiref and weboob already does the same thing.
---
 test/test_wsgi_compliance.py |  9 +++++++++
 test/wsgi_app.py             |  7 +++++++
 wsgi_intercept/__init__.py   | 36 ++++++++++++++++++++----------------
 3 files changed, 36 insertions(+), 16 deletions(-)

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*