From 792c5c4558ebd7c46744f742db6286c8b3419e52 Mon Sep 17 00:00:00 2001
From: Luiz Menezes <luiz.menezesf@gmail.com>
Date: Tue, 9 Jul 2019 18:39:36 -0300
Subject: [PATCH] Record intermediary redirect responses on aiohttp stub

---
 tests/integration/test_aiohttp.py   | 15 ++++++
 vcr/stubs/aiohttp_stubs/__init__.py | 72 ++++++++++++++++++++---------
 2 files changed, 65 insertions(+), 22 deletions(-)

diff --git a/tests/integration/test_aiohttp.py b/tests/integration/test_aiohttp.py
index 8c6810b..3664739 100644
--- a/tests/integration/test_aiohttp.py
+++ b/tests/integration/test_aiohttp.py
@@ -223,3 +223,18 @@ def test_aiohttp_test_client_json(aiohttp_client, tmpdir):
     response_json = loop.run_until_complete(response.json())
     assert response_json is None
     assert cassette.play_count == 1
+
+
+def test_redirect(aiohttp_client, tmpdir):
+    url = 'https://httpbin.org/redirect/2'
+
+    with vcr.use_cassette(str(tmpdir.join('redirect.yaml'))):
+        response, _ = get(url)
+
+    with vcr.use_cassette(str(tmpdir.join('redirect.yaml'))) as cassette:
+        cassette_response, _ = get(url)
+
+        assert cassette_response.status == response.status
+        assert len(cassette_response.history) == len(response.history)
+        assert len(cassette) == 3
+        assert cassette.play_count == 3
diff --git a/vcr/stubs/aiohttp_stubs/__init__.py b/vcr/stubs/aiohttp_stubs/__init__.py
index 0512fa2..fbe9cdf 100644
--- a/vcr/stubs/aiohttp_stubs/__init__.py
+++ b/vcr/stubs/aiohttp_stubs/__init__.py
@@ -56,6 +56,54 @@ class MockClientResponse(ClientResponse):
         return s
 
 
+def build_response(vcr_request, vcr_response, history):
+    response = MockClientResponse(vcr_request.method, URL(vcr_response.get('url')))
+    response.status = vcr_response['status']['code']
+    response._body = vcr_response['body'].get('string', b'')
+    response.reason = vcr_response['status']['message']
+    response._headers = vcr_response['headers']
+    response._history = tuple(history)
+
+    response.close()
+    return response
+
+
+def play_responses(cassette, vcr_request):
+    history = []
+    vcr_response = cassette.play_response(vcr_request)
+    response = build_response(vcr_request, vcr_response, history)
+
+    while cassette.can_play_response_for(vcr_request):
+        history.append(response)
+        vcr_response = cassette.play_response(vcr_request)
+        response = build_response(vcr_request, vcr_response, history)
+
+    return response
+
+
+async def record_response(cassette, vcr_request, response, past=False):
+    body = {} if past else {'string': (await response.read())}
+    headers = {str(key): value for key, value in response.headers.items()}
+
+    vcr_response = {
+        'status': {
+            'code': response.status,
+            'message': response.reason,
+        },
+        'headers': headers,
+        'body': body,  # NOQA: E999
+        'url': str(response.url),
+    }
+    cassette.append(vcr_request, vcr_response)
+
+
+async def record_responses(cassette, vcr_request, response):
+    for past_response in response.history:
+        await record_response(cassette, vcr_request, past_response, past=True)
+
+    await record_response(cassette, vcr_request, response)
+
+
 def vcr_request(cassette, real_request):
     @functools.wraps(real_request)
     async def new_request(self, method, url, **kwargs):
@@ -77,16 +125,7 @@ def vcr_request(cassette, real_request):
         vcr_request = Request(method, str(request_url), data, headers)
 
         if cassette.can_play_response_for(vcr_request):
-            vcr_response = cassette.play_response(vcr_request)
-
-            response = MockClientResponse(method, URL(vcr_response.get('url')))
-            response.status = vcr_response['status']['code']
-            response._body = vcr_response['body']['string']
-            response.reason = vcr_response['status']['message']
-            response._headers = vcr_response['headers']
-
-            response.close()
-            return response
+            return play_responses(cassette, vcr_request)
 
         if cassette.write_protected and cassette.filter_request(vcr_request):
             response = MockClientResponse(method, URL(url))
@@ -101,18 +140,7 @@ def vcr_request(cassette, real_request):
         log.info("{} not in cassette, sending to real server", vcr_request)
 
         response = await real_request(self, method, url, **kwargs)  # NOQA: E999
-
-        vcr_response = {
-            'status': {
-                'code': response.status,
-                'message': response.reason,
-            },
-            'headers': {str(key): value for key, value in response.headers.items()},
-            'body': {'string': (await response.read())},  # NOQA: E999
-            'url': str(response.url),
-        }
-        cassette.append(vcr_request, vcr_response)
-
+        await record_responses(cassette, vcr_request, response)
         return response
 
     return new_request