This will allow automatic connection and socket cleanup on command completion.
We also repeat the Qmp tests using the context manager. For this to be feasible, we move the test scenario to class variables and modify QmpStub to not consume its script. Signed-off-by: Apollon Oikonomopoulos <[email protected]> --- lib/hypervisor/hv_kvm/monitor.py | 7 +++ test/py/ganeti.hypervisor.hv_kvm_unittest.py | 73 +++++++++++++++++----------- 2 files changed, 51 insertions(+), 29 deletions(-) diff --git a/lib/hypervisor/hv_kvm/monitor.py b/lib/hypervisor/hv_kvm/monitor.py index 8305219..7667730 100644 --- a/lib/hypervisor/hv_kvm/monitor.py +++ b/lib/hypervisor/hv_kvm/monitor.py @@ -207,6 +207,13 @@ class QmpConnection(MonitorSocket): self._buf = "" self.supported_commands = None + def __enter__(self): + self.connect() + return self + + def __exit__(self, exc_type, exc_value, tb): + self.close() + def connect(self): """Connects to the QMP monitor. diff --git a/test/py/ganeti.hypervisor.hv_kvm_unittest.py b/test/py/ganeti.hypervisor.hv_kvm_unittest.py index aebf1f6..e4d37d5 100755 --- a/test/py/ganeti.hypervisor.hv_kvm_unittest.py +++ b/test/py/ganeti.hypervisor.hv_kvm_unittest.py @@ -86,7 +86,7 @@ class QmpStub(threading.Thread): """ threading.Thread.__init__(self) self.socket_filename = socket_filename - self.script = server_responses + self.script = server_responses[:] self.socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) self.socket.bind(self.socket_filename) @@ -168,37 +168,37 @@ class TestQmpMessage(testutils.GanetiTestCase): class TestQmp(testutils.GanetiTestCase): - def testQmp(self): - requests = [ - {"execute": "query-kvm", "arguments": []}, - {"execute": "eject", "arguments": {"device": "ide1-cd0"}}, - {"execute": "query-status", "arguments": []}, - {"execute": "query-name", "arguments": []}, - ] - - server_responses = [ - # One message, one send() - '{"return": {"enabled": true, "present": true}}\r\n', - - # Message sent using multiple send() - ['{"retur', 'n": {}}\r\n'], - - # Multiple messages sent using one send() - '{"return": [{"name": "quit"}, {"name": "eject"}]}\r\n' - '{"return": {"running": true, "singlestep": false}}\r\n', - ] - - expected_responses = [ - {"enabled": True, "present": True}, - {}, - [{"name": "quit"}, {"name": "eject"}], - {"running": True, "singlestep": False}, - ] + REQUESTS = [ + {"execute": "query-kvm", "arguments": []}, + {"execute": "eject", "arguments": {"device": "ide1-cd0"}}, + {"execute": "query-status", "arguments": []}, + {"execute": "query-name", "arguments": []}, + ] + + SERVER_RESPONSES = [ + # One message, one send() + '{"return": {"enabled": true, "present": true}}\r\n', + + # Message sent using multiple send() + ['{"retur', 'n": {}}\r\n'], + + # Multiple messages sent using one send() + '{"return": [{"name": "quit"}, {"name": "eject"}]}\r\n' + '{"return": {"running": true, "singlestep": false}}\r\n', + ] + + EXPECTED_RESPONSES = [ + {"enabled": True, "present": True}, + {}, + [{"name": "quit"}, {"name": "eject"}], + {"running": True, "singlestep": False}, + ] + def testQmp(self): # Set up the stub socket_file = tempfile.NamedTemporaryFile() os.remove(socket_file.name) - qmp_stub = QmpStub(socket_file.name, server_responses) + qmp_stub = QmpStub(socket_file.name, self.SERVER_RESPONSES) qmp_stub.start() # Set up the QMP connection @@ -206,7 +206,8 @@ class TestQmp(testutils.GanetiTestCase): qmp_connection.connect() # Format the script - for request, expected_response in zip(requests, expected_responses): + for request, expected_response in zip(self.REQUESTS, + self.EXPECTED_RESPONSES): response = qmp_connection.Execute(request["execute"], request["arguments"]) self.assertEqual(response, expected_response) @@ -218,6 +219,20 @@ class TestQmp(testutils.GanetiTestCase): qmp_connection.Execute, "unsupported-command") + def testQmpContextManager(self): + # Set up the stub + socket_file = tempfile.NamedTemporaryFile() + os.remove(socket_file.name) + qmp_stub = QmpStub(socket_file.name, self.SERVER_RESPONSES) + qmp_stub.start() + + # Test the context manager functionality + with hv_kvm.QmpConnection(socket_file.name) as qmp: + for request, expected_response in zip(self.REQUESTS, + self.EXPECTED_RESPONSES): + response = qmp.Execute(request["execute"], request["arguments"]) + self.assertEqual(response, expected_response) + class TestConsole(unittest.TestCase): def _Test(self, instance, node, group, hvparams): -- 1.9.1
