just in case if we have transport close callbacks in debug mode socket operations must fail if the socket is not in blocking mode test in nonblocking mode consume data same as basetestsockclientops but using sockrecvinto consume data consume data consume data make connection fill the buffer until sending 5 chars would block cancel a blocked socksendall receive everything that is not a space immediately make another socksendall call proactoreventloop could deliver hello so endswith is necessary after the first connect attempt before the listener is ready the socket needs time to recover to make the next connect call on linux a second retry will do on windows the waiting time is unpredictable and on freebsd the socket may never come back because it s a loopback address here we ll just retry for a few times and have to skip the test if it s not working see also https stackoverflow coma544376023316267 https lists freebsd orgpipermailfreebsdcurrent2005may049876 html retry only for this error winerror 10022 an invalid argument was supplied success http headers size is less than mtu they are sent by the first packet always strip headers http headers size is less than mtu they are sent by the first packet always strip headers happy path sock sendto returns immediately happy path sock sendto returns immediately sad path sock sendto raises blockingioerror this involves patching sock sendto to raise blockingioerror but sendto is not used by the proactor event loop make sure that we will get an unused port should always exist just in case if we have transport close callbacks in debug mode socket operations must fail if the socket is not in blocking mode test in non blocking mode consume data same as _basetest_sock_client_ops but using sock_recv_into consume data consume data consume data make connection fill the buffer until sending 5 chars would block cancel a blocked sock_sendall receive everything that is not a space immediately make another sock_sendall call proactoreventloop could deliver hello so endswith is necessary after the first connect attempt before the listener is ready the socket needs time to recover to make the next connect call on linux a second retry will do on windows the waiting time is unpredictable and on freebsd the socket may never come back because it s a loopback address here we ll just retry for a few times and have to skip the test if it s not working see also https stackoverflow com a 54437602 3316267 https lists freebsd org pipermail freebsd current 2005 may 049876 html retry only for this error winerror 10022 an invalid argument was supplied success http headers size is less than mtu they are sent by the first packet always strip headers http headers size is less than mtu they are sent by the first packet always strip headers happy path sock sendto returns immediately happy path sock sendto returns immediately sad path sock sendto raises blockingioerror this involves patching sock sendto to raise blockingioerror but sendto is not used by the proactor event loop make sure that we will get an unused port should always exist
import socket
import asyncio
import sys
import unittest

from asyncio import proactor_events
from itertools import cycle, islice
from unittest.mock import Mock
from test.test_asyncio import utils as test_utils
from test import support
from test.support import socket_helper

if socket_helper.tcp_blackhole():
    raise unittest.SkipTest('Not relevant to ProactorEventLoop')


def tearDownModule():
    asyncio.set_event_loop_policy(None)


class MyProto(asyncio.Protocol):
    connected = None
    done = None

    def __init__(self, loop=None):
        self.transport = None
        self.state = 'INITIAL'
        self.nbytes = 0
        if loop is not None:
            self.connected = loop.create_future()
            self.done = loop.create_future()

    def _assert_state(self, *expected):
        if self.state not in expected:
            raise AssertionError(f'state: {self.state!r}, expected: {expected!r}')

    def connection_made(self, transport):
        self.transport = transport
        self._assert_state('INITIAL')
        self.state = 'CONNECTED'
        if self.connected:
            self.connected.set_result(None)
        transport.write(b'GET / HTTP/1.0\r\nHost: example.com\r\n\r\n')

    def data_received(self, data):
        self._assert_state('CONNECTED')
        self.nbytes += len(data)

    def eof_received(self):
        self._assert_state('CONNECTED')
        self.state = 'EOF'

    def connection_lost(self, exc):
        self._assert_state('CONNECTED', 'EOF')
        self.state = 'CLOSED'
        if self.done:
            self.done.set_result(None)


class BaseSockTestsMixin:

    def create_event_loop(self):
        raise NotImplementedError

    def setUp(self):
        self.loop = self.create_event_loop()
        self.set_event_loop(self.loop)
        super().setUp()

    def tearDown(self):
        
        if not self.loop.is_closed():
            test_utils.run_briefly(self.loop)

        self.doCleanups()
        support.gc_collect()
        super().tearDown()

    def _basetest_sock_client_ops(self, httpd, sock):
        if not isinstance(self.loop, proactor_events.BaseProactorEventLoop):
            
            
            self.loop.set_debug(True)
            sock.setblocking(True)
            with self.assertRaises(ValueError):
                self.loop.run_until_complete(
                    self.loop.sock_connect(sock, httpd.address))
            with self.assertRaises(ValueError):
                self.loop.run_until_complete(
                    self.loop.sock_sendall(sock, b'GET / HTTP/1.0\r\n\r\n'))
            with self.assertRaises(ValueError):
                self.loop.run_until_complete(
                    self.loop.sock_recv(sock, 1024))
            with self.assertRaises(ValueError):
                self.loop.run_until_complete(
                    self.loop.sock_recv_into(sock, bytearray()))
            with self.assertRaises(ValueError):
                self.loop.run_until_complete(
                    self.loop.sock_accept(sock))

        
        sock.setblocking(False)
        self.loop.run_until_complete(
            self.loop.sock_connect(sock, httpd.address))
        self.loop.run_until_complete(
            self.loop.sock_sendall(sock, b'GET / HTTP/1.0\r\n\r\n'))
        data = self.loop.run_until_complete(
            self.loop.sock_recv(sock, 1024))
        
        self.loop.run_until_complete(
            self.loop.sock_recv(sock, 1024))
        sock.close()
        self.assertTrue(data.startswith(b'HTTP/1.0 200 OK'))

    def _basetest_sock_recv_into(self, httpd, sock):
        
        sock.setblocking(False)
        self.loop.run_until_complete(
            self.loop.sock_connect(sock, httpd.address))
        self.loop.run_until_complete(
            self.loop.sock_sendall(sock, b'GET / HTTP/1.0\r\n\r\n'))
        data = bytearray(1024)
        with memoryview(data) as buf:
            nbytes = self.loop.run_until_complete(
                self.loop.sock_recv_into(sock, buf[:1024]))
            
            self.loop.run_until_complete(
                self.loop.sock_recv_into(sock, buf[nbytes:]))
        sock.close()
        self.assertTrue(data.startswith(b'HTTP/1.0 200 OK'))

    def test_sock_client_ops(self):
        with test_utils.run_test_server() as httpd:
            sock = socket.socket()
            self._basetest_sock_client_ops(httpd, sock)
            sock = socket.socket()
            self._basetest_sock_recv_into(httpd, sock)

    async def _basetest_sock_recv_racing(self, httpd, sock):
        sock.setblocking(False)
        await self.loop.sock_connect(sock, httpd.address)

        task = asyncio.create_task(self.loop.sock_recv(sock, 1024))
        await asyncio.sleep(0)
        task.cancel()

        asyncio.create_task(
            self.loop.sock_sendall(sock, b'GET / HTTP/1.0\r\n\r\n'))
        data = await self.loop.sock_recv(sock, 1024)
        
        await self.loop.sock_recv(sock, 1024)

        self.assertTrue(data.startswith(b'HTTP/1.0 200 OK'))

    async def _basetest_sock_recv_into_racing(self, httpd, sock):
        sock.setblocking(False)
        await self.loop.sock_connect(sock, httpd.address)

        data = bytearray(1024)
        with memoryview(data) as buf:
            task = asyncio.create_task(
                self.loop.sock_recv_into(sock, buf[:1024]))
            await asyncio.sleep(0)
            task.cancel()

            task = asyncio.create_task(
                self.loop.sock_sendall(sock, b'GET / HTTP/1.0\r\n\r\n'))
            nbytes = await self.loop.sock_recv_into(sock, buf[:1024])
            
            await self.loop.sock_recv_into(sock, buf[nbytes:])
            self.assertTrue(data.startswith(b'HTTP/1.0 200 OK'))

        await task

    async def _basetest_sock_send_racing(self, listener, sock):
        listener.bind(('127.0.0.1', 0))
        listener.listen(1)

        
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 1024)
        sock.setblocking(False)
        task = asyncio.create_task(
            self.loop.sock_connect(sock, listener.getsockname()))
        await asyncio.sleep(0)
        server = listener.accept()[0]
        server.setblocking(False)

        with server:
            await task

            
            size = 8192
            while size >= 4:
                with self.assertRaises(BlockingIOError):
                    while True:
                        sock.send(b' ' * size)
                size = int(size / 2)

            
            task = asyncio.create_task(
                self.loop.sock_sendall(sock, b'hello'))
            await asyncio.sleep(0)
            task.cancel()

            
            async def recv_all():
                rv = b''
                while True:
                    buf = await self.loop.sock_recv(server, 8192)
                    if not buf:
                        return rv
                    rv += buf.strip()
            task = asyncio.create_task(recv_all())

            
            await self.loop.sock_sendall(sock, b'world')
            sock.shutdown(socket.SHUT_WR)
            data = await task
            
            self.assertTrue(data.endswith(b'world'))

    
    
    
    
    
    
    
    
    async def _basetest_sock_connect_racing(self, listener, sock):
        listener.bind(('127.0.0.1', 0))
        addr = listener.getsockname()
        sock.setblocking(False)

        task = asyncio.create_task(self.loop.sock_connect(sock, addr))
        await asyncio.sleep(0)
        task.cancel()

        listener.listen(1)

        skip_reason = "Max retries reached"
        for i in range(128):
            try:
                await self.loop.sock_connect(sock, addr)
            except ConnectionRefusedError as e:
                skip_reason = e
            except OSError as e:
                skip_reason = e

                
                
                if getattr(e, 'winerror', 0) != 10022:
                    break
            else:
                
                return

        self.skipTest(skip_reason)

    def test_sock_client_racing(self):
        with test_utils.run_test_server() as httpd:
            sock = socket.socket()
            with sock:
                self.loop.run_until_complete(asyncio.wait_for(
                    self._basetest_sock_recv_racing(httpd, sock), 10))
            sock = socket.socket()
            with sock:
                self.loop.run_until_complete(asyncio.wait_for(
                    self._basetest_sock_recv_into_racing(httpd, sock), 10))
        listener = socket.socket()
        sock = socket.socket()
        with listener, sock:
            self.loop.run_until_complete(asyncio.wait_for(
                self._basetest_sock_send_racing(listener, sock), 10))

    def test_sock_client_connect_racing(self):
        listener = socket.socket()
        sock = socket.socket()
        with listener, sock:
            self.loop.run_until_complete(asyncio.wait_for(
                self._basetest_sock_connect_racing(listener, sock), 10))

    async def _basetest_huge_content(self, address):
        sock = socket.socket()
        sock.setblocking(False)
        DATA_SIZE = 10_000_00

        chunk = b'0123456789' * (DATA_SIZE // 10)

        await self.loop.sock_connect(sock, address)
        await self.loop.sock_sendall(sock,
                                     (b'POST /loop HTTP/1.0\r\n' +
                                      b'Content-Length: %d\r\n' % DATA_SIZE +
                                      b'\r\n'))

        task = asyncio.create_task(self.loop.sock_sendall(sock, chunk))

        data = await self.loop.sock_recv(sock, DATA_SIZE)
        
        
        self.assertTrue(data.startswith(b'HTTP/1.0 200 OK'))
        while data.find(b'\r\n\r\n') == -1:
            data += await self.loop.sock_recv(sock, DATA_SIZE)
        
        headers = data[:data.index(b'\r\n\r\n') + 4]
        data = data[len(headers):]

        size = DATA_SIZE
        checker = cycle(b'0123456789')

        expected = bytes(islice(checker, len(data)))
        self.assertEqual(data, expected)
        size -= len(data)

        while True:
            data = await self.loop.sock_recv(sock, DATA_SIZE)
            if not data:
                break
            expected = bytes(islice(checker, len(data)))
            self.assertEqual(data, expected)
            size -= len(data)
        self.assertEqual(size, 0)

        await task
        sock.close()

    def test_huge_content(self):
        with test_utils.run_test_server() as httpd:
            self.loop.run_until_complete(
                self._basetest_huge_content(httpd.address))

    async def _basetest_huge_content_recvinto(self, address):
        sock = socket.socket()
        sock.setblocking(False)
        DATA_SIZE = 10_000_00

        chunk = b'0123456789' * (DATA_SIZE // 10)

        await self.loop.sock_connect(sock, address)
        await self.loop.sock_sendall(sock,
                                     (b'POST /loop HTTP/1.0\r\n' +
                                      b'Content-Length: %d\r\n' % DATA_SIZE +
                                      b'\r\n'))

        task = asyncio.create_task(self.loop.sock_sendall(sock, chunk))

        array = bytearray(DATA_SIZE)
        buf = memoryview(array)

        nbytes = await self.loop.sock_recv_into(sock, buf)
        data = bytes(buf[:nbytes])
        
        
        self.assertTrue(data.startswith(b'HTTP/1.0 200 OK'))
        while data.find(b'\r\n\r\n') == -1:
            nbytes = await self.loop.sock_recv_into(sock, buf)
            data = bytes(buf[:nbytes])
        
        headers = data[:data.index(b'\r\n\r\n') + 4]
        data = data[len(headers):]

        size = DATA_SIZE
        checker = cycle(b'0123456789')

        expected = bytes(islice(checker, len(data)))
        self.assertEqual(data, expected)
        size -= len(data)

        while True:
            nbytes = await self.loop.sock_recv_into(sock, buf)
            data = buf[:nbytes]
            if not data:
                break
            expected = bytes(islice(checker, len(data)))
            self.assertEqual(data, expected)
            size -= len(data)
        self.assertEqual(size, 0)

        await task
        sock.close()

    def test_huge_content_recvinto(self):
        with test_utils.run_test_server() as httpd:
            self.loop.run_until_complete(
                self._basetest_huge_content_recvinto(httpd.address))

    async def _basetest_datagram_recvfrom(self, server_address):
        
        data = b'\x01' * 4096
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
            sock.setblocking(False)
            await self.loop.sock_sendto(sock, data, server_address)
            received_data, from_addr = await self.loop.sock_recvfrom(
                sock, 4096)
            self.assertEqual(received_data, data)
            self.assertEqual(from_addr, server_address)

    def test_recvfrom(self):
        with test_utils.run_udp_echo_server() as server_address:
            self.loop.run_until_complete(
                self._basetest_datagram_recvfrom(server_address))

    async def _basetest_datagram_recvfrom_into(self, server_address):
        
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
            sock.setblocking(False)

            buf = bytearray(4096)
            data = b'\x01' * 4096
            await self.loop.sock_sendto(sock, data, server_address)
            num_bytes, from_addr = await self.loop.sock_recvfrom_into(
                sock, buf)
            self.assertEqual(num_bytes, 4096)
            self.assertEqual(buf, data)
            self.assertEqual(from_addr, server_address)

            buf = bytearray(8192)
            await self.loop.sock_sendto(sock, data, server_address)
            num_bytes, from_addr = await self.loop.sock_recvfrom_into(
                sock, buf, 4096)
            self.assertEqual(num_bytes, 4096)
            self.assertEqual(buf[:4096], data[:4096])
            self.assertEqual(from_addr, server_address)

    def test_recvfrom_into(self):
        with test_utils.run_udp_echo_server() as server_address:
            self.loop.run_until_complete(
                self._basetest_datagram_recvfrom_into(server_address))

    async def _basetest_datagram_sendto_blocking(self, server_address):
        
        
        
        data = b'\x01' * 4096
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
            sock.setblocking(False)
            mock_sock = Mock(sock)
            mock_sock.gettimeout = sock.gettimeout
            mock_sock.sendto.configure_mock(side_effect=BlockingIOError)
            mock_sock.fileno = sock.fileno
            self.loop.call_soon(
                lambda: setattr(mock_sock, 'sendto', sock.sendto)
            )
            await self.loop.sock_sendto(mock_sock, data, server_address)

            received_data, from_addr = await self.loop.sock_recvfrom(
                sock, 4096)
            self.assertEqual(received_data, data)
            self.assertEqual(from_addr, server_address)

    def test_sendto_blocking(self):
        if sys.platform == 'win32':
            if isinstance(self.loop, asyncio.ProactorEventLoop):
                raise unittest.SkipTest('Not relevant to ProactorEventLoop')

        with test_utils.run_udp_echo_server() as server_address:
            self.loop.run_until_complete(
                self._basetest_datagram_sendto_blocking(server_address))

    @socket_helper.skip_unless_bind_unix_socket
    def test_unix_sock_client_ops(self):
        with test_utils.run_test_unix_server() as httpd:
            sock = socket.socket(socket.AF_UNIX)
            self._basetest_sock_client_ops(httpd, sock)
            sock = socket.socket(socket.AF_UNIX)
            self._basetest_sock_recv_into(httpd, sock)

    def test_sock_client_fail(self):
        
        address = None
        try:
            s = socket.socket()
            s.bind(('127.0.0.1', 0))
            address = s.getsockname()
        finally:
            s.close()

        sock = socket.socket()
        sock.setblocking(False)
        with self.assertRaises(ConnectionRefusedError):
            self.loop.run_until_complete(
                self.loop.sock_connect(sock, address))
        sock.close()

    def test_sock_accept(self):
        listener = socket.socket()
        listener.setblocking(False)
        listener.bind(('127.0.0.1', 0))
        listener.listen(1)
        client = socket.socket()
        client.connect(listener.getsockname())

        f = self.loop.sock_accept(listener)
        conn, addr = self.loop.run_until_complete(f)
        self.assertEqual(conn.gettimeout(), 0)
        self.assertEqual(addr, client.getsockname())
        self.assertEqual(client.getpeername(), listener.getsockname())
        client.close()
        conn.close()
        listener.close()

    def test_cancel_sock_accept(self):
        listener = socket.socket()
        listener.setblocking(False)
        listener.bind(('127.0.0.1', 0))
        listener.listen(1)
        sockaddr = listener.getsockname()
        f = asyncio.wait_for(self.loop.sock_accept(listener), 0.1)
        with self.assertRaises(asyncio.TimeoutError):
            self.loop.run_until_complete(f)

        listener.close()
        client = socket.socket()
        client.setblocking(False)
        f = self.loop.sock_connect(client, sockaddr)
        with self.assertRaises(ConnectionRefusedError):
            self.loop.run_until_complete(f)

        client.close()

    def test_create_connection_sock(self):
        with test_utils.run_test_server() as httpd:
            sock = None
            infos = self.loop.run_until_complete(
                self.loop.getaddrinfo(
                    *httpd.address, type=socket.SOCK_STREAM))
            for family, type, proto, cname, address in infos:
                try:
                    sock = socket.socket(family=family, type=type, proto=proto)
                    sock.setblocking(False)
                    self.loop.run_until_complete(
                        self.loop.sock_connect(sock, address))
                except BaseException:
                    pass
                else:
                    break
            else:
                self.fail('Can not create socket.')

            f = self.loop.create_connection(
                lambda: MyProto(loop=self.loop), sock=sock)
            tr, pr = self.loop.run_until_complete(f)
            self.assertIsInstance(tr, asyncio.Transport)
            self.assertIsInstance(pr, asyncio.Protocol)
            self.loop.run_until_complete(pr.done)
            self.assertGreater(pr.nbytes, 0)
            tr.close()


if sys.platform == 'win32':

    class SelectEventLoopTests(BaseSockTestsMixin,
                               test_utils.TestCase):

        def create_event_loop(self):
            return asyncio.SelectorEventLoop()

    class ProactorEventLoopTests(BaseSockTestsMixin,
                                 test_utils.TestCase):

        def create_event_loop(self):
            return asyncio.ProactorEventLoop()

else:
    import selectors

    if hasattr(selectors, 'KqueueSelector'):
        class KqueueEventLoopTests(BaseSockTestsMixin,
                                   test_utils.TestCase):

            def create_event_loop(self):
                return asyncio.SelectorEventLoop(
                    selectors.KqueueSelector())

    if hasattr(selectors, 'EpollSelector'):
        class EPollEventLoopTests(BaseSockTestsMixin,
                                  test_utils.TestCase):

            def create_event_loop(self):
                return asyncio.SelectorEventLoop(selectors.EpollSelector())

    if hasattr(selectors, 'PollSelector'):
        class PollEventLoopTests(BaseSockTestsMixin,
                                 test_utils.TestCase):

            def create_event_loop(self):
                return asyncio.SelectorEventLoop(selectors.PollSelector())

    
    class SelectEventLoopTests(BaseSockTestsMixin,
                               test_utils.TestCase):

        def create_event_loop(self):
            return asyncio.SelectorEventLoop(selectors.SelectSelector())


if __name__ == '__main__':
    unittest.main()