import asyncio
import logging
import types

loop = asyncio.ProactorEventLoop()
asyncio.set_event_loop(loop)


logging.basicConfig(level=logging.DEBUG)


async def handle_server(reader, writer):
    await reader.readexactly(4)
    writer.write(b'test')
    await writer.drain()
    await asyncio.sleep(1)
    writer.close()
    await writer.wait_closed()
    print('closed writer')


async def run_server():
    server = await asyncio.start_server(
        handle_server, '127.0.0.1', 8888)

    addr = server.sockets[0].getsockname()
    print(f'Serving on {addr}')

    async with server:
        await server.serve_forever()


async def run_client():
    reader, writer = await asyncio.open_connection('127.0.0.1', 8888)
    proto = writer._protocol
    org_connection_lost = proto.connection_lost

    def connection_lost(self, exc):
        print('Connection lost', exc)
        org_connection_lost(exc)

    proto.connection_lost = types.MethodType(connection_lost, proto)
    writer.write(b'test')
    await reader.readexactly(4)
    await asyncio.sleep(2)

    print(reader.at_eof(), writer.transport.is_closing())  # expected both True
    writer.write(b'test')
    await writer.drain()  # This should raise

    print(reader.at_eof(), writer.transport.is_closing())
    writer.write(b'test')
    await writer.drain()


async def main():
    server = asyncio.create_task(run_server())
    try:
        await run_client()
    finally:
        server.cancel()


if __name__ == '__main__':
    asyncio.run(main())