An attempt at adding UDP support to aiosocks. Untested due to lack of server support.
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

117 lines
4.3 KiB

  1. import unittest
  2. import asyncio
  3. import aiosocks
  4. import aiohttp
  5. from unittest import mock
  6. from asyncio import coroutine
  7. from aiohttp.client_reqrep import ClientRequest
  8. from aiosocks.connector import SocksConnector
  9. class TestSocksConnector(unittest.TestCase):
  10. def setUp(self):
  11. self.loop = asyncio.new_event_loop()
  12. asyncio.set_event_loop(None)
  13. def tearDown(self):
  14. self.loop.close()
  15. def _fake_coroutine(self, return_value):
  16. def coro(*args, **kwargs):
  17. if isinstance(return_value, Exception):
  18. raise return_value
  19. return return_value
  20. return mock.Mock(side_effect=coroutine(coro))
  21. @mock.patch('aiosocks.connector.create_connection')
  22. def test_connect_proxy_ip(self, cr_conn_mock):
  23. tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol')
  24. cr_conn_mock.side_effect = \
  25. self._fake_coroutine((tr, proto)).side_effect
  26. loop_mock = mock.Mock()
  27. req = ClientRequest('GET', 'http://python.org', loop=self.loop)
  28. connector = SocksConnector(aiosocks.Socks5Addr('127.0.0.1'),
  29. None, loop=loop_mock)
  30. loop_mock.getaddrinfo = self._fake_coroutine([mock.MagicMock()])
  31. conn = self.loop.run_until_complete(connector.connect(req))
  32. self.assertTrue(loop_mock.getaddrinfo.is_called)
  33. self.assertIs(conn._transport, tr)
  34. conn.close()
  35. @mock.patch('aiosocks.connector.create_connection')
  36. def test_connect_proxy_domain(self, cr_conn_mock):
  37. tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol')
  38. cr_conn_mock.side_effect = \
  39. self._fake_coroutine((tr, proto)).side_effect
  40. loop_mock = mock.Mock()
  41. req = ClientRequest('GET', 'http://python.org', loop=self.loop)
  42. connector = SocksConnector(aiosocks.Socks5Addr('proxy.example'),
  43. None, loop=loop_mock)
  44. connector._resolve_host = self._fake_coroutine([mock.MagicMock()])
  45. conn = self.loop.run_until_complete(connector.connect(req))
  46. self.assertTrue(connector._resolve_host.is_called)
  47. self.assertEqual(connector._resolve_host.call_count, 1)
  48. self.assertIs(conn._transport, tr)
  49. conn.close()
  50. @mock.patch('aiosocks.connector.create_connection')
  51. def test_connect_locale_resolve(self, cr_conn_mock):
  52. tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol')
  53. cr_conn_mock.side_effect = \
  54. self._fake_coroutine((tr, proto)).side_effect
  55. req = ClientRequest('GET', 'http://python.org', loop=self.loop)
  56. connector = SocksConnector(aiosocks.Socks5Addr('proxy.example'),
  57. None, loop=self.loop, remote_resolve=False)
  58. connector._resolve_host = self._fake_coroutine([mock.MagicMock()])
  59. conn = self.loop.run_until_complete(connector.connect(req))
  60. self.assertTrue(connector._resolve_host.is_called)
  61. self.assertEqual(connector._resolve_host.call_count, 2)
  62. conn.close()
  63. @mock.patch('aiosocks.connector.create_connection')
  64. def test_proxy_connect_fail(self, cr_conn_mock):
  65. loop_mock = mock.Mock()
  66. cr_conn_mock.side_effect = \
  67. self._fake_coroutine(aiosocks.SocksConnectionError()).side_effect
  68. req = ClientRequest('GET', 'http://python.org', loop=self.loop)
  69. connector = SocksConnector(aiosocks.Socks5Addr('127.0.0.1'),
  70. None, loop=loop_mock)
  71. loop_mock.getaddrinfo = self._fake_coroutine([mock.MagicMock()])
  72. with self.assertRaises(aiohttp.ProxyConnectionError):
  73. self.loop.run_until_complete(connector.connect(req))
  74. @mock.patch('aiosocks.connector.create_connection')
  75. def test_proxy_negotiate_fail(self, cr_conn_mock):
  76. loop_mock = mock.Mock()
  77. cr_conn_mock.side_effect = \
  78. self._fake_coroutine(aiosocks.SocksError()).side_effect
  79. req = ClientRequest('GET', 'http://python.org', loop=self.loop)
  80. connector = SocksConnector(aiosocks.Socks5Addr('127.0.0.1'),
  81. None, loop=loop_mock)
  82. loop_mock.getaddrinfo = self._fake_coroutine([mock.MagicMock()])
  83. with self.assertRaises(aiosocks.SocksError):
  84. self.loop.run_until_complete(connector.connect(req))