diff --git a/tests/test_connector.py b/tests/test_connector.py index e614b87..e58c6fd 100644 --- a/tests/test_connector.py +++ b/tests/test_connector.py @@ -3,9 +3,9 @@ import asyncio import aiosocks import aiohttp from unittest import mock -from asyncio import coroutine from aiohttp.client_reqrep import ClientRequest from aiosocks.connector import SocksConnector +from .helpers import fake_coroutine class TestSocksConnector(unittest.TestCase): @@ -16,19 +16,11 @@ class TestSocksConnector(unittest.TestCase): def tearDown(self): self.loop.close() - def _fake_coroutine(self, return_value): - def coro(*args, **kwargs): - if isinstance(return_value, Exception): - raise return_value - return return_value - - return mock.Mock(side_effect=coroutine(coro)) - @mock.patch('aiosocks.connector.create_connection') def test_connect_proxy_ip(self, cr_conn_mock): tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol') cr_conn_mock.side_effect = \ - self._fake_coroutine((tr, proto)).side_effect + fake_coroutine((tr, proto)).side_effect loop_mock = mock.Mock() @@ -36,7 +28,7 @@ class TestSocksConnector(unittest.TestCase): connector = SocksConnector(aiosocks.Socks5Addr('127.0.0.1'), None, loop=loop_mock) - loop_mock.getaddrinfo = self._fake_coroutine([mock.MagicMock()]) + loop_mock.getaddrinfo = fake_coroutine([mock.MagicMock()]) conn = self.loop.run_until_complete(connector.connect(req)) @@ -49,14 +41,14 @@ class TestSocksConnector(unittest.TestCase): def test_connect_proxy_domain(self, cr_conn_mock): tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol') cr_conn_mock.side_effect = \ - self._fake_coroutine((tr, proto)).side_effect + fake_coroutine((tr, proto)).side_effect loop_mock = mock.Mock() req = ClientRequest('GET', 'http://python.org', loop=self.loop) connector = SocksConnector(aiosocks.Socks5Addr('proxy.example'), None, loop=loop_mock) - connector._resolve_host = self._fake_coroutine([mock.MagicMock()]) + connector._resolve_host = fake_coroutine([mock.MagicMock()]) conn = self.loop.run_until_complete(connector.connect(req)) @@ -70,13 +62,13 @@ class TestSocksConnector(unittest.TestCase): def test_connect_locale_resolve(self, cr_conn_mock): tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol') cr_conn_mock.side_effect = \ - self._fake_coroutine((tr, proto)).side_effect + fake_coroutine((tr, proto)).side_effect req = ClientRequest('GET', 'http://python.org', loop=self.loop) connector = SocksConnector(aiosocks.Socks5Addr('proxy.example'), None, loop=self.loop, remote_resolve=False) - connector._resolve_host = self._fake_coroutine([mock.MagicMock()]) + connector._resolve_host = fake_coroutine([mock.MagicMock()]) conn = self.loop.run_until_complete(connector.connect(req)) @@ -89,13 +81,13 @@ class TestSocksConnector(unittest.TestCase): def test_proxy_connect_fail(self, cr_conn_mock): loop_mock = mock.Mock() cr_conn_mock.side_effect = \ - self._fake_coroutine(aiosocks.SocksConnectionError()).side_effect + fake_coroutine(aiosocks.SocksConnectionError()).side_effect req = ClientRequest('GET', 'http://python.org', loop=self.loop) connector = SocksConnector(aiosocks.Socks5Addr('127.0.0.1'), None, loop=loop_mock) - loop_mock.getaddrinfo = self._fake_coroutine([mock.MagicMock()]) + loop_mock.getaddrinfo = fake_coroutine([mock.MagicMock()]) with self.assertRaises(aiohttp.ProxyConnectionError): self.loop.run_until_complete(connector.connect(req)) @@ -104,13 +96,13 @@ class TestSocksConnector(unittest.TestCase): def test_proxy_negotiate_fail(self, cr_conn_mock): loop_mock = mock.Mock() cr_conn_mock.side_effect = \ - self._fake_coroutine(aiosocks.SocksError()).side_effect + fake_coroutine(aiosocks.SocksError()).side_effect req = ClientRequest('GET', 'http://python.org', loop=self.loop) connector = SocksConnector(aiosocks.Socks5Addr('127.0.0.1'), None, loop=loop_mock) - loop_mock.getaddrinfo = self._fake_coroutine([mock.MagicMock()]) + loop_mock.getaddrinfo = fake_coroutine([mock.MagicMock()]) with self.assertRaises(aiosocks.SocksError): self.loop.run_until_complete(connector.connect(req))