| import asyncio | |
| from contextlib import contextmanager | |
| import os | |
| import socket | |
| from tempfile import TemporaryDirectory | |
| import avocado | |
| from qemu.qmp import ConnectError, Runstate | |
| from qemu.qmp.protocol import AsyncProtocol, StateError | |
| from qemu.qmp.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), | |
| "<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.qmp.protocol' | |
| ) | |
| def testName(self): | |
| self.proto = NullProtocol('Steve') | |
| self.assertEqual( | |
| self.proto.name, | |
| 'Steve' | |
| ) | |
| self.assertEqual( | |
| self.proto.logger.name, | |
| 'qemu.qmp.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.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='.qmp') 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='.qmp') 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) |