read_response depends on request for stream_id

3 files changed, 37 insertions, 19 deletions

diff --git a/netlib/http/http1/protocol.py b/netlib/http/http1/protocol.py
index dc33a8af..107a48d1 100644
--- a/netlib/http/http1/protocol.py
+++ b/netlib/http/http1/protocol.py
@@ -136,7 +136,7 @@ class HTTP1Protocol(semantics.ProtocolMixin):
 
     def read_response(
         self,
-        request_method,
+        request,
         body_size_limit,
         include_body=True,
     ):
@@ -175,7 +175,7 @@ class HTTP1Protocol(semantics.ProtocolMixin):
             body = self.read_http_body(
                 headers,
                 body_size_limit,
-                request_method,
+                request.method,
                 code,
                 False
             )
diff --git a/netlib/http/http2/protocol.py b/netlib/http/http2/protocol.py
index 66ce19c8..e032c2a0 100644
--- a/netlib/http/http2/protocol.py
+++ b/netlib/http/http2/protocol.py
@@ -74,7 +74,9 @@ class HTTP2Protocol(semantics.ProtocolMixin):
         if hasattr(self.tcp_handler.rfile, "reset_timestamps"):
             self.tcp_handler.rfile.reset_timestamps()
 
-        stream_id, headers, body = self._receive_transmission(include_body)
+        stream_id, headers, body = self._receive_transmission(
+            include_body=include_body,
+        )
 
         if hasattr(self.tcp_handler.rfile, "first_byte_timestamp"):
             # more accurate timestamp_start
@@ -127,7 +129,7 @@ class HTTP2Protocol(semantics.ProtocolMixin):
 
     def read_response(
         self,
-        request_method='',
+        request='',
         body_size_limit=None,
         include_body=True,
     ):
@@ -137,7 +139,10 @@ class HTTP2Protocol(semantics.ProtocolMixin):
         if hasattr(self.tcp_handler.rfile, "reset_timestamps"):
             self.tcp_handler.rfile.reset_timestamps()
 
-        stream_id, headers, body = self._receive_transmission(include_body)
+        stream_id, headers, body = self._receive_transmission(
+            stream_id=request.stream_id,
+            include_body=include_body,
+        )
 
         if hasattr(self.tcp_handler.rfile, "first_byte_timestamp"):
             # more accurate timestamp_start
@@ -145,7 +150,7 @@ class HTTP2Protocol(semantics.ProtocolMixin):
 
         if include_body:
             timestamp_end = time.time()
-        else:
+        else:  # pragma: no cover
             timestamp_end = None
 
         response = http.Response(
@@ -358,11 +363,10 @@ class HTTP2Protocol(semantics.ProtocolMixin):
 
         return [frm.to_bytes() for frm in frms]
 
-    def _receive_transmission(self, include_body=True):
+    def _receive_transmission(self, stream_id=None, include_body=True):
         # TODO: include_body is not respected
         body_expected = True
 
-        stream_id = 0
         header_block_fragment = b''
         body = b''
 
@@ -370,7 +374,7 @@ class HTTP2Protocol(semantics.ProtocolMixin):
             frm = self.read_frame()
             if (
                 (isinstance(frm, frame.HeadersFrame) or isinstance(frm, frame.ContinuationFrame)) and
-                (stream_id == 0 or frm.stream_id == stream_id)
+                (stream_id is None or frm.stream_id == stream_id)
             ):
                 stream_id = frm.stream_id
                 header_block_fragment += frm.header_block_fragment
diff --git a/netlib/http/semantics.py b/netlib/http/semantics.py
index 836af550..e388a344 100644
--- a/netlib/http/semantics.py
+++ b/netlib/http/semantics.py
@@ -337,18 +337,32 @@ class Request(object):
 
 class EmptyRequest(Request):
 
-    def __init__(self):
+    def __init__(
+        self,
+        form_in="",
+        method="",
+        scheme="",
+        host="",
+        port="",
+        path="",
+        httpversion=None,
+        headers=None,
+        body="",
+        stream_id=None
+    ):
         super(EmptyRequest, self).__init__(
-            form_in="",
-            method="",
-            scheme="",
-            host="",
-            port="",
-            path="",
-            httpversion=(0, 0),
-            headers=odict.ODictCaseless(),
-            body="",
+            form_in=form_in,
+            method=method,
+            scheme=scheme,
+            host=host,
+            port=port,
+            path=path,
+            httpversion=(httpversion or (0, 0)),
+            headers=(headers or odict.ODictCaseless()),
+            body=body,
         )
+        if stream_id:
+            self.stream_id = stream_id
 
 
 class Response(object):