Add support for running process asynchronously

labgrid/driver/sshdriver.py

@@ -9,8 +9,10 @@
 import tempfile
 import time
 from functools import cached_property
+from contextlib import contextmanager
 
 import attr
+import pexpect
 
 from ..factory import target_factory
 from ..protocol import CommandProtocol, FileTransferProtocol
@@ -22,6 +24,56 @@
 from ..util.proxy import proxymanager
 from ..util.timeout import Timeout
 from ..util.ssh import get_ssh_connect_timeout
+from ..util.readmixin import ReadMixIn
+
+
+class SSHDriverProcess(ReadMixIn):
+    def __init__(self, sub):
+        self._sub = sub
+
+    @property
+    def exitcode(self):
+        if self._sub.isalive():
+            return None
+
+        if self._sub.exitstatus is None:
+            return -self._sub.signalstatus
+        return self._sub.exitstatus
+
+    def read(self, size=1, timeout=-1):
+        return self._sub.read_nonblocking(size, timeout)
+
+    @step(args=['data'])
+    def write(self, data):
+        self._sub.write(data)
+
+    @step(result=True)
+    def poll(self):
+        return self.exitcode
+
+    @step(result=True)
+    def stop(self):
+        self._sub.close(True)
+
+    @step(args=['pattern', 'timeout'], result=True)
+    def expect(self, pattern, *, timeout=-1):
+        index = self._sub.expect(pattern, timeout=timeout)
+        return index, self._sub.before, self._sub.match, self._sub.after
+
+    @step(result=True)
+    def wait(self):
+        return self._sub.wait()
+
+    @step(args=['char'])
+    def sendcontrol(self, char):
+        self._sub.sendcontrol(char)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, typ, value, traceback):
+        self.stop()
+        return False
 
 
 @target_factory.reg_driver
@@ -45,6 +97,7 @@ def __attrs_post_init__(self):
         self._scp = self._get_tool("scp")
         self._sshfs = self._get_tool("sshfs")
         self._rsync = self._get_tool("rsync")
+        self._processes = []
 
     def _get_tool(self, name):
         if self.target.env:
@@ -80,6 +133,7 @@ def on_activate(self):
         self._start_keepalive()
 
     def on_deactivate(self):
+        assert not self._processes, "Deactivating while a command process is running is not allowed"
         try:
             self._stop_keepalive()
         finally:
@@ -242,6 +296,35 @@ def _run(self, cmd, codec="utf-8", decodeerrors="strict", timeout=None):
             stderr.pop()
         return (stdout, stderr, sub.returncode)
 
+    @Driver.check_active
+    @step(args=['cmd'])
+    @contextmanager
+    def start_process(self, cmd: str):
+        if not self._check_keepalive():
+            raise ExecutionError("Keepalive no longer running")
+
+        cmd = f"stty -echo; {cmd}"
+        complete_cmd = [self._ssh, "-o", "LogLevel=QUIET", "-x", *self.ssh_prefix,
+                        "-p", str(self.networkservice.port), "-l", self._get_username(),
+                        self.networkservice.address, "-tt", "--", '/bin/sh -c {}'.format(shlex.quote(cmd)),
+                        ]
+        self.logger.debug("Sending command: %s", complete_cmd)
+
+        try:
+            sub = pexpect.spawn(complete_cmd[0], complete_cmd[1:])
+            sub.setecho(False)
+        except:
+            raise ExecutionError(
+                "error executing command: {}".format(complete_cmd)
+            )
+
+        with SSHDriverProcess(sub) as p:
+            self._processes.append(p)
+            try:
+                yield p
+            finally:
+                self._processes.remove(p)
+
     def interact(self, cmd=None):
         assert cmd is None or isinstance(cmd, list)
 
@@ -377,7 +460,7 @@ def scp(self, *, src, dst):
                 "-o", f"ControlPath={self.control.replace('%', '%%')}",
                 src, dst,
         ]
-        
+
         if self.explicit_sftp_mode and self._scp_supports_explicit_sftp_mode():
             complete_cmd.insert(1, "-s")
         if self.explicit_scp_mode and self._scp_supports_explicit_scp_mode():

labgrid/util/readmixin.py

@@ -0,0 +1,81 @@
+from pexpect import EOF, TIMEOUT
+from .timeout import Timeout
+
+class ReadMixIn:
+    """
+    This class make it more convenient to deal with reading from devices. A
+    typical read() will either return immediately if there is buffered data, or
+    wait up to some timeout for data to appear, then return whatever happens to
+    be buffered at that time. In both of these cases, less data that requested
+    may be returned before the timeout expires, even if no EOF is encountered.
+
+    The function in this class help deal with sources like this by trying harder
+    to read data from the device.
+
+    Inspired by the function of the same name present in Rust
+    """
+    def read(self, size, timeout):
+        """
+        Stub for mixin. Must be implemented by subclass.
+        """
+        raise NotImplementedError
+
+    def read_full(self, size=-1, *, timeout=30):
+        """
+        Reads bytes until either size bytes have been read, timeout seconds
+        have elapsed, or EOF is encountered. Returns as many bytes as were read
+        until that happens.
+
+        If size is -1, read as much data as possible until either the timeout
+        or EOF is encountered.
+        """
+        t = Timeout(timeout)
+        buf = b""
+
+        while not t.expired:
+            read_size = size - len(buf) if size >= 0 else 64
+            if read_size <= 0:
+                break
+
+            try:
+                buf += self.read(read_size, t.remaining)
+            except EOF:
+                break
+            except TIMEOUT:
+                pass
+
+        return buf
+
+    def read_to_end(self, *, timeout=30):
+        """
+        Read until EOF is encountered and return the resulting data. If the
+        timeout expires before EOF, a TIMEOUT error is raised
+
+        If an exception is raised, any data read is lost
+        """
+        t = Timeout(timeout)
+        buf = b""
+
+        while True:
+            try:
+                buf += self.read(64, t.remaining)
+            except EOF:
+                break
+
+        return buf
+
+    def read_exact(self, size, *, timeout=30):
+        """
+        Read exactly size bytes. If the timeout elapses before size bytes are
+        read, a TIMEOUT error is raised. If EOF is encountered before size
+        bytes are read, an EOF error is raised
+
+        If an exception is raised, any data read is lost
+        """
+        t = Timeout(timeout)
+        buf = b""
+
+        while len(buf) < size:
+            buf += self.read(size - len(buf), t.remaining)
+
+        return buf

labgrid/util/timeout.py

@@ -7,7 +7,7 @@
 class Timeout:
     """Reperents a timeout (as a deadline)"""
     timeout = attr.ib(
-        default=120.0, validator=attr.validators.instance_of(float)
+        default=120.0, validator=attr.validators.instance_of((float, int))
     )
 
     def __attrs_post_init__(self):

tests/test_sshdriver.py

@@ -1,5 +1,6 @@
 import pytest
 import socket
+from pexpect import TIMEOUT, EOF
 
 from labgrid import Environment
 from labgrid.driver import SSHDriver, ExecutionError
@@ -218,3 +219,92 @@ def test_unix_socket_forward(ssh_localhost, tmpdir):
                 send_socket.send(test_string.encode("utf-8"))
 
                 assert client_socket.recv(16).decode("utf-8") == test_string
+
+@pytest.mark.sshusername
+def test_start_process_simple(ssh_localhost):
+    with ssh_localhost.start_process("echo Hello World") as p:
+        assert p.read_full(6, timeout=10.0) == b"Hello "
+        assert p.read(7) == b"World\r\n"
+        assert p.read_full(1, timeout=10.0) == b""
+        assert p.read_full(timeout=10.0) == b""
+        with pytest.raises(EOF):
+            p.read(10)
+
+    with ssh_localhost.start_process("echo Hello World") as p:
+        p.expect("Hello World")
+        p.expect(EOF)
+        with pytest.raises(EOF):
+            p.expect(r".")
+
+
+@pytest.mark.sshusername
+def test_start_process_timeout(ssh_localhost):
+    with ssh_localhost.start_process("cat") as p:
+        with pytest.raises(TIMEOUT):
+            p.read(100, timeout=5) == b""
+
+        with pytest.raises(TIMEOUT):
+            p.expect("Never found", timeout=5)
+
+@pytest.mark.sshusername
+def test_start_process_stream(ssh_localhost):
+    with ssh_localhost.start_process("cat") as p:
+        p.write(b"Hello World\n")
+        data = p.read_full(timeout=10.0)
+        data = data.decode("utf-8").splitlines()
+        # Only one hello world expected. Input echo is disabled
+        assert data == ["Hello World"]
+
+        assert p.poll() is None
+        p.sendcontrol('d')
+        p.expect(EOF)
+
+        assert p.poll() == 0
+
+    with ssh_localhost.start_process("cat") as p:
+        p.write(b"ABCDEF\n")
+
+        p.expect(b"ABCDEF", timeout=10.0)
+
+        assert p.poll() is None
+        p.stop()
+
+        code = p.poll()
+        assert code is not None
+        assert code != 0
+
+
+@pytest.mark.sshusername
+def test_start_process_read_to_end(ssh_localhost):
+    with ssh_localhost.start_process("echo Hello World") as p:
+        p.read_to_end() == b"Hello World"
+        p.read_to_end() == b""
+
+
+@pytest.mark.sshusername
+def test_start_process_read_to_end_timeout(ssh_localhost):
+    with ssh_localhost.start_process("cat") as p:
+        p.write(b"Hello World\n")
+        with pytest.raises(TIMEOUT):
+            p.read_to_end(timeout=5)
+
+        p.sendcontrol('d')
+        assert p.read_to_end() == b""
+
+
+@pytest.mark.sshusername
+def test_start_process_read_exact(ssh_localhost):
+    with ssh_localhost.start_process("cat") as p:
+        p.write(b"Hello World\n")
+        assert p.read_exact(6) == b"Hello "
+        assert p.read_exact(5) == b"World"
+
+        # Discard the remaining data (echoed data)
+        p.read_full(-1, timeout=1)
+
+        with pytest.raises(TIMEOUT):
+            p.read_exact(1, timeout=5)
+        p.sendcontrol('d')
+
+        with pytest.raises(EOF):
+            p.read_exact(1, timeout=5)

tests/test_timeout.py

@@ -9,12 +9,17 @@ def test_create(self):
         assert (isinstance(t, Timeout))
         t = Timeout(5.0)
         assert (isinstance(t, Timeout))
+        t = Timeout(5)
+        assert (isinstance(t, Timeout))
+
         with pytest.raises(TypeError):
-            t = Timeout(10)
+            t = Timeout('123')
         with pytest.raises(ValueError):
             t = Timeout(-1.0)
+        with pytest.raises(ValueError):
+            t = Timeout(-1)
 
-    def test_expire(self, mocker):
+    def test_expire_float(self, mocker):
         m = mocker.patch('time.monotonic')
         m.return_value = 0.0
 
@@ -29,3 +34,19 @@ def test_expire(self, mocker):
         m.return_value += 3.0
         assert t.expired
         assert t.remaining == 0.0
+
+    def test_expire_int(self, mocker):
+        m = mocker.patch('time.monotonic')
+        m.return_value = 0.0
+
+        t = Timeout(5)
+        assert not t.expired
+        assert t.remaining == 5.0
+
+        m.return_value += 3.0
+        assert not t.expired
+        assert t.remaining == 2.0
+
+        m.return_value += 3.0
+        assert t.expired
+        assert t.remaining == 0.0