ProxyConnector authorization parameters changed

aio-libs · Jul 4, 2014 · 750a095 · 750a095
1 parent b89c281
commit 750a095
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 8 deletions.
diff --git a/aiohttp/connector.py b/aiohttp/connector.py
@@ -12,7 +12,7 @@
 
 from .errors import HttpProxyError
 from .errors import ProxyConnectionError
-from .client import ClientRequest
+from .client import ClientRequest, BasicAuth
 
 
 class Connection(object):
@@ -274,14 +274,14 @@ def _create_connection(self, req, **kwargs):
 class ProxyConnector(TCPConnector):
     """Http Proxy connector."""
 
-    def __init__(self, proxy, proxy_login=None, proxy_passwd=None,
-                 *args, **kwargs):
+    def __init__(self, proxy, *args, proxy_auth=None, **kwargs):
         super().__init__(*args, **kwargs)
         self._proxy = proxy
-        self._basic_login = proxy_login
-        self._basic_passwd = proxy_passwd
+        self._proxy_auth = proxy_auth
         assert proxy.startswith('http://'), (
             "Only http proxy supported", proxy)
+        assert proxy_auth is None or isinstance(proxy_auth, BasicAuth), (
+            "proxy_auth must be None or BasicAuth() tuple", proxy_auth)
 
     @property
     def proxy(self):
@@ -292,7 +292,7 @@ def _create_connection(self, req, **kwargs):
         proxy_req = ClientRequest(
             'GET', self._proxy,
             headers={'Host': req.host},
-            auth=aiohttp.BasicAuth(self._basic_login, self._basic_passwd),
+            auth=self._proxy_auth,
             loop=self._loop)
         try:
             transport, proto = yield from super()._create_connection(proxy_req)

diff --git a/tests/test_connector.py b/tests/test_connector.py
@@ -346,10 +346,38 @@ def test_connect(self, ClientRequestMock):
 
         ClientRequestMock.assert_called_with(
             'GET', 'http://proxy.example.com',
-            auth=aiohttp.BasicAuth(None, None),
+            auth=None,
             headers={'Host': 'www.python.org'},
             loop=loop_mock)
 
+    def test_proxy_auth(self):
+        with self.assertRaisesRegex(AssertionError,
+                                    "must be None or BasicAuth"):
+            aiohttp.ProxyConnector('http://proxy.example.com',
+                                   proxy_auth=('user', 'pass'),
+                                   loop=unittest.mock.Mock())
+
+    @unittest.mock.patch('aiohttp.connector.ClientRequest')
+    def test_proxy_override_auth(self, ClientRequestMock):
+        req = ClientRequest('GET', 'http://www.python.org')
+
+        loop_mock = unittest.mock.Mock()
+        tr, proto = unittest.mock.Mock(), unittest.mock.Mock()
+        self._fake_coroutine(loop_mock.create_connection, (tr, proto))
+
+        connector = aiohttp.ProxyConnector(
+            'http://user:pass@proxy.example.com',
+            proxy_auth=aiohttp.BasicAuth(None, None),
+            loop=loop_mock)
+
+        def check_proxy_req(*args, **kw):
+            proxy_req = ClientRequest(*args, **kw)
+            self.assertNotIn('AUTHORIZATION', proxy_req.headers)
+            return proxy_req
+
+        ClientRequestMock.side_effect = check_proxy_req
+        self.loop.run_until_complete(connector._create_connection(req))
+
     def test_proxy_connection_error(self):
         connector = aiohttp.ProxyConnector('http://proxy.example.com',
                                            loop=self.loop)
@@ -374,7 +402,7 @@ def test_auth(self, ClientRequestMock):
         loop_mock = unittest.mock.Mock()
         connector = aiohttp.ProxyConnector(
             'http://proxy.example.com', loop=loop_mock,
-            proxy_login='user', proxy_passwd='pass')
+            proxy_auth=aiohttp.BasicAuth('user', 'pass'))
         connector._resolve_host = resolve_mock = unittest.mock.Mock()
         self._fake_coroutine(resolve_mock, [unittest.mock.MagicMock()])
 
@@ -392,6 +420,11 @@ def test_auth(self, ClientRequestMock):
         self.assertNotIn('AUTHORIZATION', proxy_req.headers)
         self.assertNotIn('PROXY-AUTHORIZATION', proxy_req.headers)
 
+        ClientRequestMock.assert_called_with(
+            'GET', 'http://proxy.example.com',
+            auth=aiohttp.BasicAuth('user', 'pass'),
+            loop=unittest.mock.ANY, headers=unittest.mock.ANY)
+
     @unittest.mock.patch('aiohttp.connector.ClientRequest')
     def test_auth_from_url(self, ClientRequestMock):
         proxy_req = ClientRequest('GET', 'http://user:pass@proxy.example.com')
@@ -419,6 +452,10 @@ def test_auth_from_url(self, ClientRequestMock):
         self.assertNotIn('AUTHORIZATION', proxy_req.headers)
         self.assertNotIn('PROXY-AUTHORIZATION', proxy_req.headers)
 
+        ClientRequestMock.assert_called_with(
+            'GET', 'http://user:pass@proxy.example.com',
+            auth=None, loop=unittest.mock.ANY, headers=unittest.mock.ANY)
+
     @unittest.mock.patch('aiohttp.connector.ClientRequest')
     def test_auth__not_modifying_request(self, ClientRequestMock):
         proxy_req = ClientRequest('GET', 'http://user:pass@proxy.example.com')