Use contextmanager for fake_socks_srv

tests/helpers.py

@@ -1,6 +1,8 @@
 import asyncio
+import contextlib
 import socket
 from unittest import mock
+import gc
 try:
     from asyncio import ensure_future
 except ImportError:
@@ -24,28 +26,33 @@ def find_unused_port():
     return port
 
 
-class SocksPrimitiveProtocol(asyncio.Protocol):
-    def __init__(self, write_buff):
-        self._write_buff = write_buff
-        self._transport = None
-
-    def connection_made(self, transport):
-        self._transport = transport
+@contextlib.contextmanager
+def fake_socks_srv(loop, write_buff):
+    transports = []
 
-    def data_received(self, data):
-        self._transport.write(self._write_buff)
+    class SocksPrimitiveProtocol(asyncio.Protocol):
+        _transport = None
 
-    def connection_lost(self, exc):
-        self._transport.close()
+        def connection_made(self, transport):
+            self._transport = transport
+            transports.append(transport)
 
+        def data_received(self, data):
+            self._transport.write(write_buff)
 
-@asyncio.coroutine
-def fake_socks_srv(loop, write_buff):
     port = find_unused_port()
 
     def factory():
-        return SocksPrimitiveProtocol(write_buff)
+        return SocksPrimitiveProtocol()
+
+    srv = loop.run_until_complete(
+        loop.create_server(factory, '127.0.0.1', port))
+
+    yield port
 
-    server = yield from loop.create_server(factory, '127.0.0.1', port)
-    return server, port
+    for tr in transports:
+        tr.close()
 
+    srv.close()
+    loop.run_until_complete(srv.wait_closed())
+    gc.collect()

tests/test_functional.py

@@ -1,16 +1,13 @@
 import unittest
-import aiohttp
 import aiosocks
 import asyncio
-from aiohttp import web
-from aiosocks.connector import SocksConnector
 
 try:
     from asyncio import ensure_future
 except ImportError:
     ensure_future = asyncio.async
 
-from .helpers import fake_socks_srv, find_unused_port
+from .helpers import fake_socks_srv
 
 
 class TestCreateSocks4Connection(unittest.TestCase):
@@ -22,59 +19,49 @@ class TestCreateSocks4Connection(unittest.TestCase):
         self.loop.close()
 
     def test_connect_success(self):
-        server, port = self.loop.run_until_complete(
-            fake_socks_srv(self.loop, b'\x00\x5a\x04W\x01\x01\x01\x01test')
-        )
-        addr = aiosocks.Socks4Addr('127.0.0.1', port)
-        auth = aiosocks.Socks4Auth('usr')
-        dst = ('python.org', 80)
+        with fake_socks_srv(self.loop,
+                            b'\x00\x5a\x04W\x01\x01\x01\x01test') as port:
+            addr = aiosocks.Socks4Addr('127.0.0.1', port)
+            auth = aiosocks.Socks4Auth('usr')
+            dst = ('python.org', 80)
 
-        coro = aiosocks.create_connection(
-            None, addr, auth, dst, loop=self.loop)
-        transport, protocol = self.loop.run_until_complete(coro)
+            coro = aiosocks.create_connection(
+                None, addr, auth, dst, loop=self.loop)
+            transport, protocol = self.loop.run_until_complete(coro)
 
-        self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111))
+            self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111))
 
-        data = self.loop.run_until_complete(protocol._stream_reader.read(4))
-        self.assertEqual(data, b'test')
+            data = self.loop.run_until_complete(
+                protocol._stream_reader.read(4))
+            self.assertEqual(data, b'test')
 
-        transport.close()
-        server.close()
-        self.loop.run_until_complete(server.wait_closed())
+            transport.close()
 
     def test_invalid_data(self):
-        server, port = self.loop.run_until_complete(
-            fake_socks_srv(self.loop, b'\x01\x5a\x04W\x01\x01\x01\x01')
-        )
-        addr = aiosocks.Socks4Addr('127.0.0.1', port)
-        auth = aiosocks.Socks4Auth('usr')
-        dst = ('python.org', 80)
-
-        with self.assertRaises(aiosocks.SocksError) as ct:
-            coro = aiosocks.create_connection(
-                None, addr, auth, dst, loop=self.loop)
-            self.loop.run_until_complete(coro)
-        self.assertIn('invalid data', str(ct.exception))
-
-        server.close()
-        self.loop.run_until_complete(server.wait_closed())
+        with fake_socks_srv(self.loop,
+                            b'\x01\x5a\x04W\x01\x01\x01\x01') as port:
+            addr = aiosocks.Socks4Addr('127.0.0.1', port)
+            auth = aiosocks.Socks4Auth('usr')
+            dst = ('python.org', 80)
+
+            with self.assertRaises(aiosocks.SocksError) as ct:
+                coro = aiosocks.create_connection(
+                    None, addr, auth, dst, loop=self.loop)
+                self.loop.run_until_complete(coro)
+            self.assertIn('invalid data', str(ct.exception))
 
     def test_socks_srv_error(self):
-        server, port = self.loop.run_until_complete(
-            fake_socks_srv(self.loop, b'\x00\x5b\x04W\x01\x01\x01\x01')
-        )
-        addr = aiosocks.Socks4Addr('127.0.0.1', port)
-        auth = aiosocks.Socks4Auth('usr')
-        dst = ('python.org', 80)
-
-        with self.assertRaises(aiosocks.SocksError) as ct:
-            coro = aiosocks.create_connection(
-                None, addr, auth, dst, loop=self.loop)
-            self.loop.run_until_complete(coro)
-        self.assertIn('0x5b', str(ct.exception))
+        with fake_socks_srv(self.loop,
+                            b'\x00\x5b\x04W\x01\x01\x01\x01') as port:
+            addr = aiosocks.Socks4Addr('127.0.0.1', port)
+            auth = aiosocks.Socks4Auth('usr')
+            dst = ('python.org', 80)
 
-        server.close()
-        self.loop.run_until_complete(server.wait_closed())
+            with self.assertRaises(aiosocks.SocksError) as ct:
+                coro = aiosocks.create_connection(
+                    None, addr, auth, dst, loop=self.loop)
+                self.loop.run_until_complete(coro)
+            self.assertIn('0x5b', str(ct.exception))
 
 
 class TestCreateSocks5Connect(unittest.TestCase):
@@ -86,250 +73,193 @@ class TestCreateSocks5Connect(unittest.TestCase):
         self.loop.close()
 
     def test_connect_success_anonymous(self):
-        server, port = self.loop.run_until_complete(
-            fake_socks_srv(
+        with fake_socks_srv(
                 self.loop,
-                b'\x05\x00\x05\x00\x00\x01\x01\x01\x01\x01\x04Wtest'
-            )
-        )
-        addr = aiosocks.Socks5Addr('127.0.0.1', port)
-        auth = aiosocks.Socks5Auth('usr', 'pwd')
-        dst = ('python.org', 80)
+                b'\x05\x00\x05\x00\x00\x01\x01\x01\x01\x01\x04Wtest') as port:
+            addr = aiosocks.Socks5Addr('127.0.0.1', port)
+            auth = aiosocks.Socks5Auth('usr', 'pwd')
+            dst = ('python.org', 80)
 
-        coro = aiosocks.create_connection(
-            None, addr, auth, dst, loop=self.loop)
-        transport, protocol = self.loop.run_until_complete(coro)
+            coro = aiosocks.create_connection(
+                None, addr, auth, dst, loop=self.loop)
+            transport, protocol = self.loop.run_until_complete(coro)
 
-        self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111))
+            self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111))
 
-        data = self.loop.run_until_complete(protocol._stream_reader.read(4))
-        self.assertEqual(data, b'test')
+            data = self.loop.run_until_complete(
+                protocol._stream_reader.read(4))
+            self.assertEqual(data, b'test')
 
-        transport.close()
-        server.close()
-        self.loop.run_until_complete(server.wait_closed())
+            transport.close()
 
     def test_connect_success_usr_pwd(self):
-        server, port = self.loop.run_until_complete(
-            fake_socks_srv(
+        with fake_socks_srv(
                 self.loop,
                 b'\x05\x02\x01\x00\x05\x00\x00\x01\x01\x01\x01\x01\x04Wtest'
-            )
-        )
-        addr = aiosocks.Socks5Addr('127.0.0.1', port)
-        auth = aiosocks.Socks5Auth('usr', 'pwd')
-        dst = ('python.org', 80)
-
-        coro = aiosocks.create_connection(
-            None, addr, auth, dst, loop=self.loop)
-        transport, protocol = self.loop.run_until_complete(coro)
+        ) as port:
+            addr = aiosocks.Socks5Addr('127.0.0.1', port)
+            auth = aiosocks.Socks5Auth('usr', 'pwd')
+            dst = ('python.org', 80)
 
-        self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111))
+            coro = aiosocks.create_connection(
+                None, addr, auth, dst, loop=self.loop)
+            transport, protocol = self.loop.run_until_complete(coro)
 
-        data = self.loop.run_until_complete(protocol._stream_reader.read(4))
-        self.assertEqual(data, b'test')
+            self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111))
 
-        transport.close()
-        server.close()
-        self.loop.run_until_complete(server.wait_closed())
+            data = self.loop.run_until_complete(
+                protocol._stream_reader.read(4))
+            self.assertEqual(data, b'test')
+            transport.close()
 
     def test_auth_ver_err(self):
-        server, port = self.loop.run_until_complete(
-            fake_socks_srv(self.loop, b'\x04\x02')
-        )
-        addr = aiosocks.Socks5Addr('127.0.0.1', port)
-        auth = aiosocks.Socks5Auth('usr', 'pwd')
-        dst = ('python.org', 80)
-
-        with self.assertRaises(aiosocks.SocksError) as ct:
-            coro = aiosocks.create_connection(
-                None, addr, auth, dst, loop=self.loop)
-            self.loop.run_until_complete(coro)
-        self.assertIn('invalid version', str(ct.exception))
+        with fake_socks_srv(self.loop, b'\x04\x02') as port:
+            addr = aiosocks.Socks5Addr('127.0.0.1', port)
+            auth = aiosocks.Socks5Auth('usr', 'pwd')
+            dst = ('python.org', 80)
 
-        server.close()
-        self.loop.run_until_complete(server.wait_closed())
+            with self.assertRaises(aiosocks.SocksError) as ct:
+                coro = aiosocks.create_connection(
+                    None, addr, auth, dst, loop=self.loop)
+                self.loop.run_until_complete(coro)
+            self.assertIn('invalid version', str(ct.exception))
 
     def test_auth_method_rejected(self):
-        server, port = self.loop.run_until_complete(
-            fake_socks_srv(self.loop, b'\x05\xFF')
-        )
-        addr = aiosocks.Socks5Addr('127.0.0.1', port)
-        auth = aiosocks.Socks5Auth('usr', 'pwd')
-        dst = ('python.org', 80)
-
-        with self.assertRaises(aiosocks.SocksError) as ct:
-            coro = aiosocks.create_connection(
-                None, addr, auth, dst, loop=self.loop)
-            self.loop.run_until_complete(coro)
-        self.assertIn('authentication methods were rejected',
-                      str(ct.exception))
-
-        server.close()
-        self.loop.run_until_complete(server.wait_closed())
+        with fake_socks_srv(self.loop, b'\x05\xFF') as port:
+            addr = aiosocks.Socks5Addr('127.0.0.1', port)
+            auth = aiosocks.Socks5Auth('usr', 'pwd')
+            dst = ('python.org', 80)
+
+            with self.assertRaises(aiosocks.SocksError) as ct:
+                coro = aiosocks.create_connection(
+                    None, addr, auth, dst, loop=self.loop)
+                self.loop.run_until_complete(coro)
+            self.assertIn('authentication methods were rejected',
+                          str(ct.exception))
 
     def test_auth_status_invalid(self):
-        server, port = self.loop.run_until_complete(
-            fake_socks_srv(self.loop, b'\x05\xF0')
-        )
-        addr = aiosocks.Socks5Addr('127.0.0.1', port)
-        auth = aiosocks.Socks5Auth('usr', 'pwd')
-        dst = ('python.org', 80)
-
-        with self.assertRaises(aiosocks.SocksError) as ct:
-            coro = aiosocks.create_connection(
-                None, addr, auth, dst, loop=self.loop)
-            self.loop.run_until_complete(coro)
-        self.assertIn('invalid data', str(ct.exception))
+        with fake_socks_srv(self.loop, b'\x05\xF0') as port:
+            addr = aiosocks.Socks5Addr('127.0.0.1', port)
+            auth = aiosocks.Socks5Auth('usr', 'pwd')
+            dst = ('python.org', 80)
 
-        server.close()
-        self.loop.run_until_complete(server.wait_closed())
+            with self.assertRaises(aiosocks.SocksError) as ct:
+                coro = aiosocks.create_connection(
+                    None, addr, auth, dst, loop=self.loop)
+                self.loop.run_until_complete(coro)
+            self.assertIn('invalid data', str(ct.exception))
 
     def test_auth_status_invalid2(self):
-        server, port = self.loop.run_until_complete(
-            fake_socks_srv(self.loop, b'\x05\x02\x02\x00')
-        )
-        addr = aiosocks.Socks5Addr('127.0.0.1', port)
-        auth = aiosocks.Socks5Auth('usr', 'pwd')
-        dst = ('python.org', 80)
-
-        with self.assertRaises(aiosocks.SocksError) as ct:
-            coro = aiosocks.create_connection(
-                None, addr, auth, dst, loop=self.loop)
-            self.loop.run_until_complete(coro)
-        self.assertIn('invalid data', str(ct.exception))
+        with fake_socks_srv(self.loop, b'\x05\x02\x02\x00') as port:
+            addr = aiosocks.Socks5Addr('127.0.0.1', port)
+            auth = aiosocks.Socks5Auth('usr', 'pwd')
+            dst = ('python.org', 80)
 
-        server.close()
-        self.loop.run_until_complete(server.wait_closed())
+            with self.assertRaises(aiosocks.SocksError) as ct:
+                coro = aiosocks.create_connection(
+                    None, addr, auth, dst, loop=self.loop)
+                self.loop.run_until_complete(coro)
+            self.assertIn('invalid data', str(ct.exception))
 
     def test_auth_failed(self):
-        server, port = self.loop.run_until_complete(
-            fake_socks_srv(self.loop, b'\x05\x02\x01\x01')
-        )
-        addr = aiosocks.Socks5Addr('127.0.0.1', port)
-        auth = aiosocks.Socks5Auth('usr', 'pwd')
-        dst = ('python.org', 80)
-
-        with self.assertRaises(aiosocks.SocksError) as ct:
-            coro = aiosocks.create_connection(
-                None, addr, auth, dst, loop=self.loop)
-            self.loop.run_until_complete(coro)
-        self.assertIn('authentication failed', str(ct.exception))
+        with fake_socks_srv(self.loop, b'\x05\x02\x01\x01') as port:
+            addr = aiosocks.Socks5Addr('127.0.0.1', port)
+            auth = aiosocks.Socks5Auth('usr', 'pwd')
+            dst = ('python.org', 80)
 
-        server.close()
-        self.loop.run_until_complete(server.wait_closed())
+            with self.assertRaises(aiosocks.SocksError) as ct:
+                coro = aiosocks.create_connection(
+                    None, addr, auth, dst, loop=self.loop)
+                self.loop.run_until_complete(coro)
+            self.assertIn('authentication failed', str(ct.exception))
 
     def test_cmd_ver_err(self):
-        server, port = self.loop.run_until_complete(
-            fake_socks_srv(self.loop, b'\x05\x02\x01\x00\x04\x00\x00')
-        )
-        addr = aiosocks.Socks5Addr('127.0.0.1', port)
-        auth = aiosocks.Socks5Auth('usr', 'pwd')
-        dst = ('python.org', 80)
-
-        with self.assertRaises(aiosocks.SocksError) as ct:
-            coro = aiosocks.create_connection(
-                None, addr, auth, dst, loop=self.loop)
-            self.loop.run_until_complete(coro)
-        self.assertIn('invalid version', str(ct.exception))
-
-        server.close()
-        self.loop.run_until_complete(server.wait_closed())
+        with fake_socks_srv(self.loop,
+                            b'\x05\x02\x01\x00\x04\x00\x00') as port:
+            addr = aiosocks.Socks5Addr('127.0.0.1', port)
+            auth = aiosocks.Socks5Auth('usr', 'pwd')
+            dst = ('python.org', 80)
+
+            with self.assertRaises(aiosocks.SocksError) as ct:
+                coro = aiosocks.create_connection(
+                    None, addr, auth, dst, loop=self.loop)
+                self.loop.run_until_complete(coro)
+            self.assertIn('invalid version', str(ct.exception))
 
     def test_cmd_not_granted(self):
-        server, port = self.loop.run_until_complete(
-            fake_socks_srv(self.loop, b'\x05\x02\x01\x00\x05\x01\x00')
-        )
-        addr = aiosocks.Socks5Addr('127.0.0.1', port)
-        auth = aiosocks.Socks5Auth('usr', 'pwd')
-        dst = ('python.org', 80)
-
-        with self.assertRaises(aiosocks.SocksError) as ct:
-            coro = aiosocks.create_connection(
-                None, addr, auth, dst, loop=self.loop)
-            self.loop.run_until_complete(coro)
-        self.assertIn('General SOCKS server failure', str(ct.exception))
-
-        server.close()
-        self.loop.run_until_complete(server.wait_closed())
+        with fake_socks_srv(self.loop,
+                            b'\x05\x02\x01\x00\x05\x01\x00') as port:
+            addr = aiosocks.Socks5Addr('127.0.0.1', port)
+            auth = aiosocks.Socks5Auth('usr', 'pwd')
+            dst = ('python.org', 80)
+
+            with self.assertRaises(aiosocks.SocksError) as ct:
+                coro = aiosocks.create_connection(
+                    None, addr, auth, dst, loop=self.loop)
+                self.loop.run_until_complete(coro)
+            self.assertIn('General SOCKS server failure', str(ct.exception))
 
     def test_invalid_address_type(self):
-        server, port = self.loop.run_until_complete(
-            fake_socks_srv(self.loop, b'\x05\x02\x01\x00\x05\x00\x00\xFF')
-        )
-        addr = aiosocks.Socks5Addr('127.0.0.1', port)
-        auth = aiosocks.Socks5Auth('usr', 'pwd')
-        dst = ('python.org', 80)
-
-        with self.assertRaises(aiosocks.SocksError) as ct:
-            coro = aiosocks.create_connection(
-                None, addr, auth, dst, loop=self.loop)
-            self.loop.run_until_complete(coro)
-        self.assertIn('invalid data', str(ct.exception))
-
-        server.close()
-        self.loop.run_until_complete(server.wait_closed())
+        with fake_socks_srv(self.loop,
+                            b'\x05\x02\x01\x00\x05\x00\x00\xFF') as port:
+            addr = aiosocks.Socks5Addr('127.0.0.1', port)
+            auth = aiosocks.Socks5Auth('usr', 'pwd')
+            dst = ('python.org', 80)
+
+            with self.assertRaises(aiosocks.SocksError) as ct:
+                coro = aiosocks.create_connection(
+                    None, addr, auth, dst, loop=self.loop)
+                self.loop.run_until_complete(coro)
+            self.assertIn('invalid data', str(ct.exception))
 
     def test_atype_ipv4(self):
-        server, port = self.loop.run_until_complete(
-            fake_socks_srv(
+        with fake_socks_srv(
                 self.loop,
                 b'\x05\x02\x01\x00\x05\x00\x00\x01\x01\x01\x01\x01\x04W'
-            )
-        )
-        addr = aiosocks.Socks5Addr('127.0.0.1', port)
-        auth = aiosocks.Socks5Auth('usr', 'pwd')
-        dst = ('python.org', 80)
+        ) as port:
+            addr = aiosocks.Socks5Addr('127.0.0.1', port)
+            auth = aiosocks.Socks5Auth('usr', 'pwd')
+            dst = ('python.org', 80)
 
-        coro = aiosocks.create_connection(
-            None, addr, auth, dst, loop=self.loop)
-        transport, protocol = self.loop.run_until_complete(coro)
+            coro = aiosocks.create_connection(
+                None, addr, auth, dst, loop=self.loop)
+            transport, protocol = self.loop.run_until_complete(coro)
 
-        self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111))
+            self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111))
 
-        transport.close()
-        server.close()
-        self.loop.run_until_complete(server.wait_closed())
+            transport.close()
 
     def test_atype_ipv6(self):
-        server, port = self.loop.run_until_complete(
-            fake_socks_srv(
+        with fake_socks_srv(
                 self.loop,
                 b'\x05\x02\x01\x00\x05\x00\x00\x04\x00\x00\x00\x00'
-                b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x11\x04W')
-        )
-        addr = aiosocks.Socks5Addr('127.0.0.1', port)
-        auth = aiosocks.Socks5Auth('usr', 'pwd')
-        dst = ('python.org', 80)
+                b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x11\x04W'
+        ) as port:
+            addr = aiosocks.Socks5Addr('127.0.0.1', port)
+            auth = aiosocks.Socks5Auth('usr', 'pwd')
+            dst = ('python.org', 80)
 
-        coro = aiosocks.create_connection(
-            None, addr, auth, dst, loop=self.loop)
-        transport, protocol = self.loop.run_until_complete(coro)
+            coro = aiosocks.create_connection(
+                None, addr, auth, dst, loop=self.loop)
+            transport, protocol = self.loop.run_until_complete(coro)
 
-        self.assertEqual(protocol.proxy_sockname, ('::111', 1111))
+            self.assertEqual(protocol.proxy_sockname, ('::111', 1111))
 
-        transport.close()
-        server.close()
-        self.loop.run_until_complete(server.wait_closed())
+            transport.close()
 
     def test_atype_domain(self):
-        server, port = self.loop.run_until_complete(
-            fake_socks_srv(
+        with fake_socks_srv(
                 self.loop,
                 b'\x05\x02\x01\x00\x05\x00\x00\x03\x0apython.org\x04W'
-            )
-        )
-        addr = aiosocks.Socks5Addr('127.0.0.1', port)
-        auth = aiosocks.Socks5Auth('usr', 'pwd')
-        dst = ('python.org', 80)
-
-        coro = aiosocks.create_connection(
-            None, addr, auth, dst, loop=self.loop)
-        transport, protocol = self.loop.run_until_complete(coro)
+        ) as port:
+            addr = aiosocks.Socks5Addr('127.0.0.1', port)
+            auth = aiosocks.Socks5Auth('usr', 'pwd')
+            dst = ('python.org', 80)
 
-        self.assertEqual(protocol.proxy_sockname, (b'python.org', 1111))
+            coro = aiosocks.create_connection(
+                None, addr, auth, dst, loop=self.loop)
+            transport, protocol = self.loop.run_until_complete(coro)
 
-        transport.close()
-        server.close()
-        self.loop.run_until_complete(server.wait_closed())
+            self.assertEqual(protocol.proxy_sockname, (b'python.org', 1111))
 
+            transport.close()