From df0e4ebd7e116fe28eaf80db34aa75352317d9a7 Mon Sep 17 00:00:00 2001
From: Jamie Lennox <jamielennox@redhat.com>
Date: Wed, 20 Aug 2014 06:11:51 +1000
Subject: [PATCH] Correctly set adapter to response.connection

The response.connection should be the adapter that the request went
through. This is important in some situations (requests-kerberos) where
the connection is reused by the auth plugin to send additional requests.

Change-Id: I87bf996e2edbfb29eba9b4000d976a49cbf6c8cc
Closes-Bug: #1358903
---
 requests_mock/adapter.py             |  1 +
 requests_mock/exceptions.py          |  4 ++++
 requests_mock/response.py            | 10 +++++++++-
 requests_mock/tests/test_adapter.py  | 28 ++++++++++++++++++++++++++++
 requests_mock/tests/test_response.py | 11 +++++++++++
 5 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/requests_mock/adapter.py b/requests_mock/adapter.py
index 83d6c2b..d5aca3c 100644
--- a/requests_mock/adapter.py
+++ b/requests_mock/adapter.py
@@ -215,6 +215,7 @@ class Adapter(BaseAdapter, _RequestHistoryTracker):
         for matcher in reversed(self._matchers):
             resp = matcher(request)
             if resp is not None:
+                resp.connection = self
                 return resp
 
         raise exceptions.NoMockAddress(request)
diff --git a/requests_mock/exceptions.py b/requests_mock/exceptions.py
index 69e3f8a..97a9820 100644
--- a/requests_mock/exceptions.py
+++ b/requests_mock/exceptions.py
@@ -24,3 +24,7 @@ class NoMockAddress(MockException):
     def __str__(self):
         return "No mock address: %s %s" % (self.request.method,
                                            self.request.url)
+
+
+class InvalidRequest(MockException):
+    """This call cannot be made under a mocked environment"""
diff --git a/requests_mock/response.py b/requests_mock/response.py
index 4b707ed..06acb95 100644
--- a/requests_mock/response.py
+++ b/requests_mock/response.py
@@ -17,6 +17,7 @@ from requests.packages.urllib3.response import HTTPResponse
 import six
 
 from requests_mock import compat
+from requests_mock import exceptions
 
 _BODY_ARGS = frozenset(['raw', 'body', 'content', 'text', 'json'])
 _HTTP_ARGS = frozenset(['status_code', 'reason', 'headers'])
@@ -43,6 +44,11 @@ def _check_body_arguments(**kwargs):
 class _FakeConnection(object):
     """An object that can mock the necessary parts of a socket interface."""
 
+    def send(self, request, **kwargs):
+        msg = 'This response was created without a connection. You are ' \
+              'therefore unable to make a request directly on that connection.'
+        raise exceptions.InvalidRequest(msg)
+
     def close(self):
         pass
 
@@ -62,6 +68,8 @@ def create_response(request, **kwargs):
     :param dict headers: A dictionary object containing headers that are
         returned upon a successful match.
     """
+    connection = kwargs.pop('connection', _FakeConnection())
+
     _check_body_arguments(**kwargs)
 
     raw = kwargs.pop('raw', None)
@@ -93,7 +101,7 @@ def create_response(request, **kwargs):
                            original_response=compat._fake_http_response)
 
     response = _http_adapter.build_response(request, raw)
-    response.connection = _FakeConnection()
+    response.connection = connection
     response.encoding = encoding
     return response
 
diff --git a/requests_mock/tests/test_adapter.py b/requests_mock/tests/test_adapter.py
index 5d66c7f..5171112 100644
--- a/requests_mock/tests/test_adapter.py
+++ b/requests_mock/tests/test_adapter.py
@@ -367,3 +367,31 @@ class SessionAdapterTests(base.TestCase):
         resp = self.session.get(good)
 
         self.assertEqual('good', resp.text)
+
+    def test_adapter_is_connection(self):
+        url = '%s://test.url' % self.PREFIX
+        text = 'text'
+        self.adapter.register_uri('GET', url, text=text)
+        resp = self.session.get(url)
+
+        self.assertEqual(text, resp.text)
+        self.assertIs(self.adapter, resp.connection)
+
+    def test_send_to_connection(self):
+        url1 = '%s://test1.url/' % self.PREFIX
+        url2 = '%s://test2.url/' % self.PREFIX
+
+        text1 = 'text1'
+        text2 = 'text2'
+
+        self.adapter.register_uri('GET', url1, text=text1)
+        self.adapter.register_uri('GET', url2, text=text2)
+
+        req = requests.Request(method='GET', url=url2)
+        req = self.session.prepare_request(req)
+
+        resp1 = self.session.get(url1)
+        self.assertEqual(text1, resp1.text)
+
+        resp2 = resp1.connection.send(req)
+        self.assertEqual(text2, resp2.text)
diff --git a/requests_mock/tests/test_response.py b/requests_mock/tests/test_response.py
index 804e83b..a571e7f 100644
--- a/requests_mock/tests/test_response.py
+++ b/requests_mock/tests/test_response.py
@@ -13,6 +13,7 @@
 import six
 
 from requests_mock import adapter
+from requests_mock import exceptions
 from requests_mock import response
 from requests_mock.tests import base
 
@@ -66,3 +67,13 @@ class ResponseTests(base.TestCase):
         self.assertEqual(value, resp.text)
         self.assertIsInstance(resp.text, six.string_types)
         self.assertIsInstance(resp.content, six.binary_type)
+
+    def test_setting_connection(self):
+        conn = object()
+        resp = self.create_response(connection=conn)
+        self.assertIs(conn, resp.connection)
+
+    def test_send_from_no_connection(self):
+        resp = self.create_response()
+        self.assertRaises(exceptions.InvalidRequest,
+                          resp.connection.send, self.request)