Merge remote branch 'cros/upstream' into tempbranch

client/tests/kvm/kvm_monitor.py

Index: client/tests/kvm/kvm_monitor.py
diff --git a/client/tests/kvm/kvm_monitor.py b/client/tests/kvm/kvm_monitor.py
index 8440835fd2487aaf0ddc810a7b114a10c08cb0e1..7e6b594bc242d10574f42a0aac1e869348ccd339 100644
--- a/client/tests/kvm/kvm_monitor.py
+++ b/client/tests/kvm/kvm_monitor.py
@@ -4,7 +4,7 @@ Interfaces to the QEMU monitor.
 @copyright: 2008-2010 Red Hat Inc.
 """
 
-import socket, time, threading, logging
+import socket, time, threading, logging, select
 import kvm_utils
 try:
     import json
@@ -21,7 +21,7 @@ class MonitorConnectError(MonitorError):
     pass
 
 
-class MonitorSendError(MonitorError):
+class MonitorSocketError(MonitorError):
     pass
 
 
@@ -38,7 +38,15 @@ class MonitorNotSupportedError(MonitorError):
 
 
 class QMPCmdError(MonitorError):
-    pass
+    def __init__(self, cmd, qmp_args, data):
+        MonitorError.__init__(self, cmd, qmp_args, data)
+        self.cmd = cmd
+        self.qmp_args = qmp_args
+        self.data = data
+
+    def __str__(self):
+        return ("QMP command %r failed (arguments: %r, error message: %r)" %
+                (self.cmd, self.qmp_args, self.data))
 
 
 class Monitor:
@@ -58,7 +66,6 @@ class Monitor:
         self.filename = filename
         self._lock = threading.RLock()
         self._socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-        self._socket.setblocking(False)
 
         try:
             self._socket.connect(filename)
@@ -102,13 +109,19 @@ class Monitor:
         return False
 
 
+    def _data_available(self, timeout=0):
+        timeout = max(0, timeout)
+        return bool(select.select([self._socket], [], [], timeout)[0])
+
+
     def _recvall(self):
         s = ""
-        while True:
+        while self._data_available():
             try:
                 data = self._socket.recv(1024)
-            except socket.error:
-                break
+            except socket.error, (errno, msg):
+                raise MonitorSocketError("Could not receive data from monitor "
+                                         "(%s)" % msg)
             if not data:
                 break
             s += data
@@ -130,7 +143,7 @@ class HumanMonitor(Monitor):
                 suppress_exceptions is False
         @raise MonitorProtocolError: Raised if the initial (qemu) prompt isn't
                 found and suppress_exceptions is False
-        @note: Other exceptions may be raised.  See _get_command_output's
+        @note: Other exceptions may be raised.  See cmd()'s
                 docstring.
         """
         try:
@@ -146,7 +159,7 @@ class HumanMonitor(Monitor):
                                            "Output so far: %r" % o)
 
             # Save the output of 'help' for future use
-            self._help_str = self._get_command_output("help")
+            self._help_str = self.cmd("help")
 
         except MonitorError, e:
             if suppress_exceptions:
@@ -158,46 +171,47 @@ class HumanMonitor(Monitor):
     # Private methods
 
     def _read_up_to_qemu_prompt(self, timeout=20):
-        o = ""
+        s = ""
         end_time = time.time() + timeout
-        while time.time() < end_time:
+        while self._data_available(end_time - time.time()):
+            data = self._recvall()
+            if not data:
+                break
+            s += data
             try:
-                data = self._socket.recv(1024)
-                if not data:
-                    break
-                o += data
-                if o.splitlines()[-1].split()[-1] == "(qemu)":
-                    return True, "\n".join(o.splitlines()[:-1])
-            except (socket.error, IndexError):
-                time.sleep(0.01)
-        return False, "\n".join(o.splitlines())
+                if s.splitlines()[-1].split()[-1] == "(qemu)":
+                    return True, "\n".join(s.splitlines()[:-1])
+            except IndexError:
+                continue
+        return False, "\n".join(s.splitlines())
 
 
-    def _send_command(self, command):
+    def _send(self, cmd):
         """
         Send a command without waiting for output.
 
-        @param command: Command to send
-        @return: True if successful, False otherwise
+        @param cmd: Command to send
         @raise MonitorLockError: Raised if the lock cannot be acquired
-        @raise MonitorSendError: Raised if the command cannot be sent
+        @raise MonitorSocketError: Raised if a socket error occurs
         """
         if not self._acquire_lock(20):
             raise MonitorLockError("Could not acquire exclusive lock to send "
-                                   "monitor command '%s'" % command)
+                                   "monitor command '%s'" % cmd)
 
         try:
             try:
-                self._socket.sendall(command + "\n")
-            except socket.error:
-                raise MonitorSendError("Could not send monitor command '%s'" %
-                                       command)
+                self._socket.sendall(cmd + "\n")
+            except socket.error, (errno, msg):
+                raise MonitorSocketError("Could not send monitor command '%s' "
+                                         "(%s)" % (cmd, msg))
 
         finally:
             self._lock.release()
 
 
-    def _get_command_output(self, command, timeout=20):
+    # Public methods
+
+    def cmd(self, command, timeout=20):
         """
         Send command to the monitor.
 
@@ -205,7 +219,7 @@ class HumanMonitor(Monitor):
         @param timeout: Time duration to wait for the (qemu) prompt to return
         @return: Output received from the monitor
         @raise MonitorLockError: Raised if the lock cannot be acquired
-        @raise MonitorSendError: Raised if the command cannot be sent
+        @raise MonitorSocketError: Raised if a socket error occurs
         @raise MonitorProtocolError: Raised if the (qemu) prompt cannot be
                 found after sending the command
         """
@@ -217,7 +231,7 @@ class HumanMonitor(Monitor):
             # Read any data that might be available
             self._recvall()
             # Send command
-            self._send_command(command)
+            self._send(command)
             # Read output
             s, o = self._read_up_to_qemu_prompt(timeout)
             # Remove command echo from output
@@ -234,8 +248,6 @@ class HumanMonitor(Monitor):
             self._lock.release()
 
 
-    # Public methods
-
     def is_responsive(self):
         """
         Make sure the monitor is responsive by sending a command.
@@ -243,7 +255,7 @@ class HumanMonitor(Monitor):
         @return: True if responsive, False otherwise
         """
         try:
-            self._get_command_output("help")
+            self.cmd("info status")
             return True
         except MonitorError:
             return False
@@ -252,39 +264,22 @@ class HumanMonitor(Monitor):
     # Command wrappers
     # Notes:
     # - All of the following commands raise exceptions in a similar manner to
-    #   cmd() and _get_command_output().
+    #   cmd().
     # - A command wrapper should use self._help_str if it requires information
     #   about the monitor's capabilities.
 
-    def cmd(self, command, timeout=20):
-        """
-        Send a simple command with no parameters and return its output.
-        Should only be used for commands that take no parameters and are
-        implemented under the same name for both the human and QMP monitors.
-
-        @param command: Command to send
-        @param timeout: Time duration to wait for (qemu) prompt after command
-        @return: The output of the command
-        @raise MonitorLockError: Raised if the lock cannot be acquired
-        @raise MonitorSendError: Raised if the command cannot be sent
-        @raise MonitorProtocolError: Raised if the (qemu) prompt cannot be
-                found after sending the command
-        """
-        return self._get_command_output(command, timeout)
-
-
     def quit(self):
         """
         Send "quit" without waiting for output.
         """
-        self._send_command("quit")
+        self._send("quit")
 
 
     def info(self, what):
         """
         Request info about something and return the output.
         """
-        return self._get_command_output("info %s" % what)
+        return self.cmd("info %s" % what)
 
 
     def query(self, what):
@@ -301,7 +296,7 @@ class HumanMonitor(Monitor):
         @param filename: Location for the screendump
         @return: The command's output
         """
-        return self._get_command_output("screendump %s" % filename)
+        return self.cmd("screendump %s" % filename)
 
 
     def migrate(self, uri, full_copy=False, incremental_copy=False, wait=False):
@@ -323,7 +318,7 @@ class HumanMonitor(Monitor):
         if incremental_copy:
             cmd += " -i"
         cmd += " %s" % uri
-        return self._get_command_output(cmd)
+        return self.cmd(cmd)
 
 
     def migrate_set_speed(self, value):
@@ -333,7 +328,7 @@ class HumanMonitor(Monitor):
         @param value: Speed in bytes/sec
         @return: The command's output
         """
-        return self._get_command_output("migrate_set_speed %s" % value)
+        return self.cmd("migrate_set_speed %s" % value)
 
 
     def sendkey(self, keystr, hold_time=1):
@@ -344,7 +339,7 @@ class HumanMonitor(Monitor):
         @param hold_time: Hold time in ms (should normally stay 1 ms)
         @return: The command's output
         """
-        return self._get_command_output("sendkey %s %s" % (keystr, hold_time))
+        return self.cmd("sendkey %s %s" % (keystr, hold_time))
 
 
     def mouse_move(self, dx, dy):
@@ -355,7 +350,7 @@ class HumanMonitor(Monitor):
         @param dy: Y amount
         @return: The command's output
         """
-        return self._get_command_output("mouse_move %d %d" % (dx, dy))
+        return self.cmd("mouse_move %d %d" % (dx, dy))
 
 
     def mouse_button(self, state):
@@ -365,7 +360,7 @@ class HumanMonitor(Monitor):
         @param state: Button state (1=L, 2=M, 4=R)
         @return: The command's output
         """
-        return self._get_command_output("mouse_button %d" % state)
+        return self.cmd("mouse_button %d" % state)
 
 
 class QMPMonitor(Monitor):
@@ -387,7 +382,7 @@ class QMPMonitor(Monitor):
         @raise MonitorNotSupportedError: Raised if json isn't available and
                 suppress_exceptions is False
         @note: Other exceptions may be raised if the qmp_capabilities command
-                fails.  See _get_command_output's docstring.
+                fails.  See cmd()'s docstring.
         """
         try:
             Monitor.__init__(self, name, filename)
@@ -408,7 +403,7 @@ class QMPMonitor(Monitor):
             while time.time() < end_time:
                 for obj in self._read_objects():
                     if "QMP" in obj:
-                        self._greeting = obj["QMP"]
+                        self._greeting = obj
                         break
                 if self._greeting:
                     break
@@ -417,7 +412,7 @@ class QMPMonitor(Monitor):
                 raise MonitorProtocolError("No QMP greeting message received")
 
             # Issue qmp_capabilities
-            self._get_command_output("qmp_capabilities")
+            self.cmd("qmp_capabilities")
 
         except MonitorError, e:
             if suppress_exceptions:
@@ -439,7 +434,7 @@ class QMPMonitor(Monitor):
 
     def _read_objects(self, timeout=5):
         """
-        Read lines from monitor and try to decode them.
+        Read lines from the monitor and try to decode them.
         Stop when all available lines have been successfully decoded, or when
         timeout expires.  If any decoded objects are asynchronous events, store
         them in self._events.  Return all decoded objects.
@@ -447,99 +442,178 @@ class QMPMonitor(Monitor):
         @param timeout: Time to wait for all lines to decode successfully
         @return: A list of objects
         """
+        if not self._data_available():
+            return []
         s = ""
-        objs = []
         end_time = time.time() + timeout
-        while time.time() < end_time:
+        while self._data_available(end_time - time.time()):
             s += self._recvall()
+            # Make sure all lines are decodable
             for line in s.splitlines():
-                if not line:
-                    continue
-                try:
-                    obj = json.loads(line)
-                except:
-                    # Found an incomplete or broken line -- keep reading
-                    break
-                objs += [obj]
+                if line:
+                    try:
+                        json.loads(line)
+                    except:
+                        # Found an incomplete or broken line -- keep reading
+                        break
             else:
                 # All lines are OK -- stop reading
                 break
-            time.sleep(0.1)
+        # Decode all decodable lines
+        objs = []
+        for line in s.splitlines():
+            try:
+                objs += [json.loads(line)]
+            except:
+                pass
         # Keep track of asynchronous events
         self._events += [obj for obj in objs if "event" in obj]
         return objs
 
 
-    def _send_command(self, cmd, args=None, id=None):
+    def _send(self, data):
+        """
+        Send raw data without waiting for response.
+
+        @param data: Data to send
+        @raise MonitorSocketError: Raised if a socket error occurs
+        """
+        try:
+            self._socket.sendall(data)
+        except socket.error, (errno, msg):
+            raise MonitorSocketError("Could not send data: %r (%s)" %
+                                     (data, msg))
+
+
+    def _get_response(self, id=None, timeout=20):
         """
-        Send command without waiting for response.
+        Read a response from the QMP monitor.
+
+        @param id: If not None, look for a response with this id
+        @param timeout: Time duration to wait for response
+        @return: The response dict, or None if none was found
+        """
+        end_time = time.time() + timeout
+        while self._data_available(end_time - time.time()):
+            for obj in self._read_objects():
+                if isinstance(obj, dict):
+                    if id is not None and obj.get("id") != id:
+                        continue
+                    if "return" in obj or "error" in obj:
+                        return obj
+
+
+    # Public methods
+
+    def cmd(self, cmd, args=None, timeout=20):
+        """
+        Send a QMP monitor command and return the response.
+
+        Note: an id is automatically assigned to the command and the response
+        is checked for the presence of the same id.
 
         @param cmd: Command to send
         @param args: A dict containing command arguments, or None
+        @param timeout: Time duration to wait for response
+        @return: The response received
         @raise MonitorLockError: Raised if the lock cannot be acquired
-        @raise MonitorSendError: Raised if the command cannot be sent
+        @raise MonitorSocketError: Raised if a socket error occurs
+        @raise MonitorProtocolError: Raised if no response is received
+        @raise QMPCmdError: Raised if the response is an error message
+                (the exception's args are (cmd, args, data) where data is the
+                error data)
         """
         if not self._acquire_lock(20):
             raise MonitorLockError("Could not acquire exclusive lock to send "
                                    "QMP command '%s'" % cmd)
 
         try:
-            cmdobj = self._build_cmd(cmd, args, id)
-            try:
-                self._socket.sendall(json.dumps(cmdobj) + "\n")
-            except socket.error:
-                raise MonitorSendError("Could not send QMP command '%s'" % cmd)
+            # Read any data that might be available
+            self._read_objects()
+            # Send command
+            id = kvm_utils.generate_random_string(8)
+            self._send(json.dumps(self._build_cmd(cmd, args, id)) + "\n")
+            # Read response
+            r = self._get_response(id, timeout)
+            if r is None:
+                raise MonitorProtocolError("Received no response to QMP "
+                                           "command '%s', or received a "
+                                           "response with an incorrect id"
+                                           % cmd)
+            if "return" in r:
+                return r["return"]
+            if "error" in r:
+                raise QMPCmdError(cmd, args, r["error"])
 
         finally:
             self._lock.release()
 
 
-    def _get_command_output(self, cmd, args=None, timeout=20):
+    def cmd_raw(self, data, timeout=20):
         """
-        Send monitor command and wait for response.
+        Send a raw string to the QMP monitor and return the response.
+        Unlike cmd(), return the raw response dict without performing any
+        checks on it.
 
-        @param cmd: Command to send
-        @param args: A dict containing command arguments, or None
+        @param data: The data to send
         @param timeout: Time duration to wait for response
         @return: The response received
         @raise MonitorLockError: Raised if the lock cannot be acquired
-        @raise MonitorSendError: Raised if the command cannot be sent
+        @raise MonitorSocketError: Raised if a socket error occurs
         @raise MonitorProtocolError: Raised if no response is received
-        @raise QMPCmdError: Raised if the response is an error message
-                (the exception's args are (msg, data) where msg is a string and
-                data is the error data)
         """
         if not self._acquire_lock(20):
             raise MonitorLockError("Could not acquire exclusive lock to send "
-                                   "QMP command '%s'" % cmd)
+                                   "data: %r" % data)
 
         try:
-            # Read any data that might be available
             self._read_objects()
-            # Send command
-            id = kvm_utils.generate_random_string(8)
-            self._send_command(cmd, args, id)
-            # Read response
-            end_time = time.time() + timeout
-            while time.time() < end_time:
-                for obj in self._read_objects():
-                    if isinstance(obj, dict) and obj.get("id") == id:
-                        if "return" in obj:
-                            return obj["return"]
-                        elif "error" in obj:
-                            raise QMPCmdError("QMP command '%s' failed" % cmd,
-                                              obj["error"])
-                time.sleep(0.1)
-            # No response found
-            raise MonitorProtocolError("Received no response to QMP command "
-                                       "'%s', or received a response with an "
-                                       "incorrect id" % cmd)
+            self._send(data)
+            r = self._get_response(None, timeout)
+            if r is None:
+                raise MonitorProtocolError("Received no response to data: %r" %
+                                           data)
+            return r
 
         finally:
             self._lock.release()
 
 
-    # Public methods
+    def cmd_obj(self, obj, timeout=20):
+        """
+        Transform a Python object to JSON, send the resulting string to the QMP
+        monitor, and return the response.
+        Unlike cmd(), return the raw response dict without performing any
+        checks on it.
+
+        @param obj: The object to send
+        @param timeout: Time duration to wait for response
+        @return: The response received
+        @raise MonitorLockError: Raised if the lock cannot be acquired
+        @raise MonitorSocketError: Raised if a socket error occurs
+        @raise MonitorProtocolError: Raised if no response is received
+        """
+        return self.cmd_raw(json.dumps(obj) + "\n")
+
+
+    def cmd_qmp(self, cmd, args=None, id=None, timeout=20):
+        """
+        Build a QMP command from the passed arguments, send it to the monitor
+        and return the response.
+        Unlike cmd(), return the raw response dict without performing any
+        checks on it.
+
+        @param cmd: Command to send
+        @param args: A dict containing command arguments, or None
+        @param id:  An id for the command, or None
+        @param timeout: Time duration to wait for response
+        @return: The response received
+        @raise MonitorLockError: Raised if the lock cannot be acquired
+        @raise MonitorSocketError: Raised if a socket error occurs
+        @raise MonitorProtocolError: Raised if no response is received
+        """
+        return self.cmd_obj(self._build_cmd(cmd, args, id), timeout)
+
 
     def is_responsive(self):
         """
@@ -548,7 +622,7 @@ class QMPMonitor(Monitor):
         @return: True if responsive, False otherwise
         """
         try:
-            self._get_command_output("query-version")
+            self.cmd("query-status")
             return True
         except MonitorError:
             return False
@@ -597,38 +671,29 @@ class QMPMonitor(Monitor):
         self._lock.release()
 
 
-    # Command wrappers
-    # Note: all of the following functions raise exceptions in a similar manner
-    # to cmd() and _get_command_output().
-
-    def cmd(self, command, timeout=20):
+    def get_greeting(self):
         """
-        Send a simple command with no parameters and return its output.
-        Should only be used for commands that take no parameters and are
-        implemented under the same name for both the human and QMP monitors.
-
-        @param command: Command to send
-        @param timeout: Time duration to wait for response
-        @return: The response to the command
-        @raise MonitorLockError: Raised if the lock cannot be acquired
-        @raise MonitorSendError: Raised if the command cannot be sent
-        @raise MonitorProtocolError: Raised if no response is received
+        Return QMP greeting message.
         """
-        return self._get_command_output(command, timeout=timeout)
+        return self._greeting
 
 
+    # Command wrappers
+    # Note: all of the following functions raise exceptions in a similar manner
+    # to cmd().
+
     def quit(self):
         """
         Send "quit" and return the response.
         """
-        return self._get_command_output("quit")
+        return self.cmd("quit")
 
 
     def info(self, what):
         """
         Request info about something and return the response.
         """
-        return self._get_command_output("query-%s" % what)
+        return self.cmd("query-%s" % what)
 
 
     def query(self, what):
@@ -646,7 +711,7 @@ class QMPMonitor(Monitor):
         @return: The response to the command
         """
         args = {"filename": filename}
-        return self._get_command_output("screendump", args)
+        return self.cmd("screendump", args)
 
 
     def migrate(self, uri, full_copy=False, incremental_copy=False, wait=False):
@@ -662,7 +727,7 @@ class QMPMonitor(Monitor):
         args = {"uri": uri,
                 "blk": full_copy,
                 "inc": incremental_copy}
-        return self._get_command_output("migrate", args)
+        return self.cmd("migrate", args)
 
 
     def migrate_set_speed(self, value):
@@ -673,4 +738,4 @@ class QMPMonitor(Monitor):
         @return: The response to the command
         """
         args = {"value": value}
-        return self._get_command_output("migrate_set_speed", args)
+        return self.cmd("migrate_set_speed", args)