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)
