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_start_server(self, address, ssl=None): if self.fake_session: self._accepted = asyncio.Event() self._set_state(Runstate.CONNECTING) await asyncio.sleep(0) else: await super()._do_start_server(address, ssl) async def _do_accept(self): if self.fake_session: self._accepted = None else: await super()._do_accept() async def _do_connect(self, address, ssl=None): if self.fake_session: self._set_state(Runstate.CONNECTING) await asyncio.sleep(0) else: 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), "" ) 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), "" ) 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.start_server_and_accept(('example.com', 1)) elif family == 'UNIX': await self.proto.start_server_and_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.start_server_and_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.start_server_and_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.start_server_and_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.start_server_and_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.start_server_and_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.start_server_and_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.start_server_and_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.start_server_and_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.start_server_and_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.start_server_and_accept(sock)) # give the server a chance to start listening [...] await asyncio.sleep(0) await self.proto.connect(sock)