qemu/python/tests/protocol.py

import asyncio
from contextlib import contextmanager
import os
import socket
from tempfile import TemporaryDirectory

import avocado

from qemu.aqmp import ConnectError, Runstate
from qemu.aqmp.protocol import AsyncProtocol, StateError
from qemu.aqmp.util import asyncio_run, create_task


class NullProtocol(AsyncProtocol[None]):
    """
    NullProtocol is a test mockup of an AsyncProtocol implementation.

    It adds a fake_session instance variable that enables a code path
    that bypasses the actual connection logic, but still allows the
    reader/writers to start.

    Because the message type is defined as None, an asyncio.Event named
    'trigger_input' is created that prohibits the reader from
    incessantly being able to yield None; this event can be poked to
    simulate an incoming message.

    For testing symmetry with do_recv, an interface is added to "send" a
    Null message.

    For testing purposes, a "simulate_disconnection" method is also
    added which allows us to trigger a bottom half disconnect without
    injecting any real errors into the reader/writer loops; in essence
    it performs exactly half of what disconnect() normally does.
    """
    def __init__(self, name=None):
        self.fake_session = False
        self.trigger_input: asyncio.Event
        super().__init__(name)

    async def _establish_session(self):
        self.trigger_input = asyncio.Event()
        await super()._establish_session()

    async def _do_accept(self, address, ssl=None):
        if not self.fake_session:
            await super()._do_accept(address, ssl)

    async def _do_connect(self, address, ssl=None):
        if not self.fake_session:
            await super()._do_connect(address, ssl)

    async def _do_recv(self) -> None:
        await self.trigger_input.wait()
        self.trigger_input.clear()

    def _do_send(self, msg: None) -> None:
        pass

    async def send_msg(self) -> None:
        await self._outgoing.put(None)

    async def simulate_disconnect(self) -> None:
        """
        Simulates a bottom-half disconnect.

        This method schedules a disconnection but does not wait for it
        to complete. This is used to put the loop into the DISCONNECTING
        state without fully quiescing it back to IDLE. This is normally
        something you cannot coax AsyncProtocol to do on purpose, but it
        will be similar to what happens with an unhandled Exception in
        the reader/writer.

        Under normal circumstances, the library design requires you to
        await on disconnect(), which awaits the disconnect task and
        returns bottom half errors as a pre-condition to allowing the
        loop to return back to IDLE.
        """
        self._schedule_disconnect()


class LineProtocol(AsyncProtocol[str]):
    def __init__(self, name=None):
        super().__init__(name)
        self.rx_history = []

    async def _do_recv(self) -> str:
        raw = await self._readline()
        msg = raw.decode()
        self.rx_history.append(msg)
        return msg

    def _do_send(self, msg: str) -> None:
        assert self._writer is not None
        self._writer.write(msg.encode() + b'\n')

    async def send_msg(self, msg: str) -> None:
        await self._outgoing.put(msg)


def run_as_task(coro, allow_cancellation=False):
    """
    Run a given coroutine as a task.

    Optionally, wrap it in a try..except block that allows this
    coroutine to be canceled gracefully.
    """
    async def _runner():
        try:
            await coro
        except asyncio.CancelledError:
            if allow_cancellation:
                return
            raise
    return create_task(_runner())


@contextmanager
def jammed_socket():
    """
    Opens up a random unused TCP port on localhost, then jams it.
    """
    socks = []

    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind(('127.0.0.1', 0))
        sock.listen(1)
        address = sock.getsockname()

        socks.append(sock)

        # I don't *fully* understand why, but it takes *two* un-accepted
        # connections to start jamming the socket.
        for _ in range(2):
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.connect(address)
            socks.append(sock)

        yield address

    finally:
        for sock in socks:
            sock.close()


class Smoke(avocado.Test):

    def setUp(self):
        self.proto = NullProtocol()

    def test__repr__(self):
        self.assertEqual(
            repr(self.proto),
            "<NullProtocol runstate=IDLE>"
        )

    def testRunstate(self):
        self.assertEqual(
            self.proto.runstate,
            Runstate.IDLE
        )

    def testDefaultName(self):
        self.assertEqual(
            self.proto.name,
            None
        )

    def testLogger(self):
        self.assertEqual(
            self.proto.logger.name,
            'qemu.aqmp.protocol'
        )

    def testName(self):
        self.proto = NullProtocol('Steve')

        self.assertEqual(
            self.proto.name,
            'Steve'
        )

        self.assertEqual(
            self.proto.logger.name,
            'qemu.aqmp.protocol.Steve'
        )

        self.assertEqual(
            repr(self.proto),
            "<NullProtocol name='Steve' runstate=IDLE>"
        )


class TestBase(avocado.Test):

    def setUp(self):
        self.proto = NullProtocol(type(self).__name__)
        self.assertEqual(self.proto.runstate, Runstate.IDLE)
        self.runstate_watcher = None

    def tearDown(self):
        self.assertEqual(self.proto.runstate, Runstate.IDLE)

    async def _asyncSetUp(self):
        pass

    async def _asyncTearDown(self):
        if self.runstate_watcher:
            await self.runstate_watcher

    @staticmethod
    def async_test(async_test_method):
        """
        Decorator; adds SetUp and TearDown to async tests.
        """
        async def _wrapper(self, *args, **kwargs):
            loop = asyncio.get_event_loop()
            loop.set_debug(True)

            await self._asyncSetUp()
            await async_test_method(self, *args, **kwargs)
            await self._asyncTearDown()

        return _wrapper

    # Definitions

    # The states we expect a "bad" connect/accept attempt to transition through
    BAD_CONNECTION_STATES = (
        Runstate.CONNECTING,
        Runstate.DISCONNECTING,
        Runstate.IDLE,
    )

    # The states we expect a "good" session to transition through
    GOOD_CONNECTION_STATES = (
        Runstate.CONNECTING,
        Runstate.RUNNING,
        Runstate.DISCONNECTING,
        Runstate.IDLE,
    )

    # Helpers

    async def _watch_runstates(self, *states):
        """
        This launches a task alongside (most) tests below to confirm that
        the sequence of runstate changes that occur is exactly as
        anticipated.
        """
        async def _watcher():
            for state in states:
                new_state = await self.proto.runstate_changed()
                self.assertEqual(
                    new_state,
                    state,
                    msg=f"Expected state '{state.name}'",
                )

        self.runstate_watcher = create_task(_watcher())
        # Kick the loop and force the task to block on the event.
        await asyncio.sleep(0)


class State(TestBase):

    @TestBase.async_test
    async def testSuperfluousDisconnect(self):
        """
        Test calling disconnect() while already disconnected.
        """
        await self._watch_runstates(
            Runstate.DISCONNECTING,
            Runstate.IDLE,
        )
        await self.proto.disconnect()


class Connect(TestBase):
    """
    Tests primarily related to calling Connect().
    """
    async def _bad_connection(self, family: str):
        assert family in ('INET', 'UNIX')

        if family == 'INET':
            await self.proto.connect(('127.0.0.1', 0))
        elif family == 'UNIX':
            await self.proto.connect('/dev/null')

    async def _hanging_connection(self):
        with jammed_socket() as addr:
            await self.proto.connect(addr)

    async def _bad_connection_test(self, family: str):
        await self._watch_runstates(*self.BAD_CONNECTION_STATES)

        with self.assertRaises(ConnectError) as context:
            await self._bad_connection(family)

        self.assertIsInstance(context.exception.exc, OSError)
        self.assertEqual(
            context.exception.error_message,
            "Failed to establish connection"
        )

    @TestBase.async_test
    async def testBadINET(self):
        """
        Test an immediately rejected call to an IP target.
        """
        await self._bad_connection_test('INET')

    @TestBase.async_test
    async def testBadUNIX(self):
        """
        Test an immediately rejected call to a UNIX socket target.
        """
        await self._bad_connection_test('UNIX')

    @TestBase.async_test
    async def testCancellation(self):
        """
        Test what happens when a connection attempt is aborted.
        """
        # Note that accept() cannot be cancelled outright, as it isn't a task.
        # However, we can wrap it in a task and cancel *that*.
        await self._watch_runstates(*self.BAD_CONNECTION_STATES)
        task = run_as_task(self._hanging_connection(), allow_cancellation=True)

        state = await self.proto.runstate_changed()
        self.assertEqual(state, Runstate.CONNECTING)

        # This is insider baseball, but the connection attempt has
        # yielded *just* before the actual connection attempt, so kick
        # the loop to make sure it's truly wedged.
        await asyncio.sleep(0)

        task.cancel()
        await task

    @TestBase.async_test
    async def testTimeout(self):
        """
        Test what happens when a connection attempt times out.
        """
        await self._watch_runstates(*self.BAD_CONNECTION_STATES)
        task = run_as_task(self._hanging_connection())

        # More insider baseball: to improve the speed of this test while
        # guaranteeing that the connection even gets a chance to start,
        # verify that the connection hangs *first*, then await the
        # result of the task with a nearly-zero timeout.

        state = await self.proto.runstate_changed()
        self.assertEqual(state, Runstate.CONNECTING)
        await asyncio.sleep(0)

        with self.assertRaises(asyncio.TimeoutError):
            await asyncio.wait_for(task, timeout=0)

    @TestBase.async_test
    async def testRequire(self):
        """
        Test what happens when a connection attempt is made while CONNECTING.
        """
        await self._watch_runstates(*self.BAD_CONNECTION_STATES)
        task = run_as_task(self._hanging_connection(), allow_cancellation=True)

        state = await self.proto.runstate_changed()
        self.assertEqual(state, Runstate.CONNECTING)

        with self.assertRaises(StateError) as context:
            await self._bad_connection('UNIX')

        self.assertEqual(
            context.exception.error_message,
            "NullProtocol is currently connecting."
        )
        self.assertEqual(context.exception.state, Runstate.CONNECTING)
        self.assertEqual(context.exception.required, Runstate.IDLE)

        task.cancel()
        await task

    @TestBase.async_test
    async def testImplicitRunstateInit(self):
        """
        Test what happens if we do not wait on the runstate event until
        AFTER a connection is made, i.e., connect()/accept() themselves
        initialize the runstate event. All of the above tests force the
        initialization by waiting on the runstate *first*.
        """
        task = run_as_task(self._hanging_connection(), allow_cancellation=True)

        # Kick the loop to coerce the state change
        await asyncio.sleep(0)
        assert self.proto.runstate == Runstate.CONNECTING

        # We already missed the transition to CONNECTING
        await self._watch_runstates(Runstate.DISCONNECTING, Runstate.IDLE)

        task.cancel()
        await task


class Accept(Connect):
    """
    All of the same tests as Connect, but using the accept() interface.
    """
    async def _bad_connection(self, family: str):
        assert family in ('INET', 'UNIX')

        if family == 'INET':
            await self.proto.accept(('example.com', 1))
        elif family == 'UNIX':
            await self.proto.accept('/dev/null')

    async def _hanging_connection(self):
        with TemporaryDirectory(suffix='.aqmp') as tmpdir:
            sock = os.path.join(tmpdir, type(self.proto).__name__ + ".sock")
            await self.proto.accept(sock)


class FakeSession(TestBase):

    def setUp(self):
        super().setUp()
        self.proto.fake_session = True

    async def _asyncSetUp(self):
        await super()._asyncSetUp()
        await self._watch_runstates(*self.GOOD_CONNECTION_STATES)

    async def _asyncTearDown(self):
        await self.proto.disconnect()
        await super()._asyncTearDown()

    ####

    @TestBase.async_test
    async def testFakeConnect(self):

        """Test the full state lifecycle (via connect) with a no-op session."""
        await self.proto.connect('/not/a/real/path')
        self.assertEqual(self.proto.runstate, Runstate.RUNNING)

    @TestBase.async_test
    async def testFakeAccept(self):
        """Test the full state lifecycle (via accept) with a no-op session."""
        await self.proto.accept('/not/a/real/path')
        self.assertEqual(self.proto.runstate, Runstate.RUNNING)

    @TestBase.async_test
    async def testFakeRecv(self):
        """Test receiving a fake/null message."""
        await self.proto.accept('/not/a/real/path')

        logname = self.proto.logger.name
        with self.assertLogs(logname, level='DEBUG') as context:
            self.proto.trigger_input.set()
            self.proto.trigger_input.clear()
            await asyncio.sleep(0)  # Kick reader.

        self.assertEqual(
            context.output,
            [f"DEBUG:{logname}:<-- None"],
        )

    @TestBase.async_test
    async def testFakeSend(self):
        """Test sending a fake/null message."""
        await self.proto.accept('/not/a/real/path')

        logname = self.proto.logger.name
        with self.assertLogs(logname, level='DEBUG') as context:
            # Cheat: Send a Null message to nobody.
            await self.proto.send_msg()
            # Kick writer; awaiting on a queue.put isn't sufficient to yield.
            await asyncio.sleep(0)

        self.assertEqual(
            context.output,
            [f"DEBUG:{logname}:--> None"],
        )

    async def _prod_session_api(
            self,
            current_state: Runstate,
            error_message: str,
            accept: bool = True
    ):
        with self.assertRaises(StateError) as context:
            if accept:
                await self.proto.accept('/not/a/real/path')
            else:
                await self.proto.connect('/not/a/real/path')

        self.assertEqual(context.exception.error_message, error_message)
        self.assertEqual(context.exception.state, current_state)
        self.assertEqual(context.exception.required, Runstate.IDLE)

    @TestBase.async_test
    async def testAcceptRequireRunning(self):
        """Test that accept() cannot be called when Runstate=RUNNING"""
        await self.proto.accept('/not/a/real/path')

        await self._prod_session_api(
            Runstate.RUNNING,
            "NullProtocol is already connected and running.",
            accept=True,
        )

    @TestBase.async_test
    async def testConnectRequireRunning(self):
        """Test that connect() cannot be called when Runstate=RUNNING"""
        await self.proto.accept('/not/a/real/path')

        await self._prod_session_api(
            Runstate.RUNNING,
            "NullProtocol is already connected and running.",
            accept=False,
        )

    @TestBase.async_test
    async def testAcceptRequireDisconnecting(self):
        """Test that accept() cannot be called when Runstate=DISCONNECTING"""
        await self.proto.accept('/not/a/real/path')

        # Cheat: force a disconnect.
        await self.proto.simulate_disconnect()

        await self._prod_session_api(
            Runstate.DISCONNECTING,
            ("NullProtocol is disconnecting."
             " Call disconnect() to return to IDLE state."),
            accept=True,
        )

    @TestBase.async_test
    async def testConnectRequireDisconnecting(self):
        """Test that connect() cannot be called when Runstate=DISCONNECTING"""
        await self.proto.accept('/not/a/real/path')

        # Cheat: force a disconnect.
        await self.proto.simulate_disconnect()

        await self._prod_session_api(
            Runstate.DISCONNECTING,
            ("NullProtocol is disconnecting."
             " Call disconnect() to return to IDLE state."),
            accept=False,
        )


class SimpleSession(TestBase):

    def setUp(self):
        super().setUp()
        self.server = LineProtocol(type(self).__name__ + '-server')

    async def _asyncSetUp(self):
        await super()._asyncSetUp()
        await self._watch_runstates(*self.GOOD_CONNECTION_STATES)

    async def _asyncTearDown(self):
        await self.proto.disconnect()
        try:
            await self.server.disconnect()
        except EOFError:
            pass
        await super()._asyncTearDown()

    @TestBase.async_test
    async def testSmoke(self):
        with TemporaryDirectory(suffix='.aqmp') as tmpdir:
            sock = os.path.join(tmpdir, type(self.proto).__name__ + ".sock")
            server_task = create_task(self.server.accept(sock))

            # give the server a chance to start listening [...]
            await asyncio.sleep(0)
            await self.proto.connect(sock)