From e18aa9580d3e89beaf7406e30d77d8901ac18122 Mon Sep 17 00:00:00 2001
From: Tim Head <betatim@gmail.com>
Date: Tue, 30 Apr 2019 13:34:40 +0200
Subject: [PATCH] Let the callback result set the URL of a response

---
 README.rst                 |  2 +-
 aioresponses/core.py       |  6 +++++-
 tests/test_aioresponses.py | 15 +++++++++++++++
 3 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/README.rst b/README.rst
index 7be1c2b..3bf4349 100644
--- a/README.rst
+++ b/README.rst
@@ -190,7 +190,7 @@ for convenience use *payload* argument to mock out json response. Example below.
         # will throw an exception.
 
 
-**aioresponses allows to use callbacks to provide dynamic responses**
+**aioresponses allows you to use callbacks to provide dynamic responses**
 
 .. code:: python
 
diff --git a/aioresponses/core.py b/aioresponses/core.py
index 8e502be..6c46842 100644
--- a/aioresponses/core.py
+++ b/aioresponses/core.py
@@ -36,7 +36,8 @@ def __init__(self, method: str = hdrs.METH_GET,
                  payload: Dict = None,
                  headers: Dict = None,
                  response_class: 'ClientResponse' = None,
-                 reason: Optional[str] = None):
+                 reason: Optional[str] = None,
+                 url: Optional[URL] = None):
         self.method = method
         self.status = status
         self.body = body
@@ -45,6 +46,7 @@ def __init__(self, method: str = hdrs.METH_GET,
         self.headers = headers
         self.response_class = response_class
         self.reason = reason
+        self.url = url
 
 
 class RequestMatch(object):
@@ -166,6 +168,8 @@ async def build_response(
             return self.exception
         if callable(self.callback):
             result = self.callback(url, **kwargs)
+            if result.url is not None:
+                url = result.url
         else:
             result = None
         result = self if result is None else result
diff --git a/tests/test_aioresponses.py b/tests/test_aioresponses.py
index 5652faf..a469613 100644
--- a/tests/test_aioresponses.py
+++ b/tests/test_aioresponses.py
@@ -328,3 +328,18 @@ def callback(url, **kwargs):
         response = self.run_async(self.request(self.url))
         data = self.run_async(response.read())
         assert data == body
+
+    @aioresponses()
+    def test_callback_setting_url(self, m):
+        body = b'New body'
+
+        def callback(url, **kwargs):
+            self.assertEqual(str(url), self.url)
+            self.assertEqual(kwargs, {'allow_redirects': True})
+            return CallbackResult(body=body, url=url.with_query(foo='bar'))
+
+        m.get(self.url, callback=callback)
+        response = self.run_async(self.request(self.url))
+        data = self.run_async(response.read())
+        assert response.url == URL(self.url).with_query(foo='bar')
+        assert data == body