fix: preserve headers through retry test decorator (#334)

testbench/common.py

@@ -641,6 +641,12 @@ def _extract_data(data):
     return data
 
 
+def _extract_headers(response):
+    if isinstance(response, flask.Response):
+        return response.headers
+    return dict()
+
+
 def __get_streamer_response_fn(database, method, conn, test_id, limit=4, chunk_size=4):
     def response_handler(data):
         def streamer():
@@ -666,7 +672,7 @@ def streamer():
                         "Injected 'broken stream' fault", 500
                     )
 
-        return flask.Response(streamer())
+        return flask.Response(streamer(), headers=_extract_headers(data))
 
     return response_handler
 

tests/test_testbench_retry.py

@@ -425,6 +425,7 @@ def test_retry_test_return_broken_stream_after_bytes(self):
             query_string={"alt": "media"},
             headers={"x-retry-test-id": id},
         )
+        self.assertIn("x-goog-generation", response.headers)
         with self.assertRaises(testbench.error.RestException) as ex:
             _ = len(response.data)
         self.assertIn("broken stream", ex.exception.msg)