From c98b25ed098d02d031e6f369e4ba14634bf27e18 Mon Sep 17 00:00:00 2001
From: Petr Viktorin <pviktori@redhat.com>
Date: Jan 12 2018 17:42:05 +0000
Subject: [PATCH 1/5] Fix some issues with long-running programs and encodings


With OpenSSH, the stdin_text feature relied on the executed program
to read as much input as it needed, end exit. This made it
unusable with long-running programs like nc (netcat).
Fix this by passing the input as escaped string and piping it through
echo.

Based heavily on a patch by Aleksei Slaikovskii:
https://pagure.io/python-pytest-multihost/pull-request/9

---

diff --git a/pytest_multihost/host.py b/pytest_multihost/host.py
index a095a43..ec26496 100644
--- a/pytest_multihost/host.py
+++ b/pytest_multihost/host.py
@@ -233,6 +233,11 @@ class BaseHost(object):
         if self.command_prelude:
             command.stdin.write(self.command_prelude)
 
+        if stdin_text:
+            command.stdin.write("echo -e '")
+            command.stdin.write(stdin_text.replace("'", r"'\''"))
+            command.stdin.write("' | ")
+
         if isinstance(argv, basestring):
             # Run a shell command given as a string
             command.stdin.write('(')
@@ -244,9 +249,7 @@ class BaseHost(object):
                 command.stdin.write(shell_quote(arg))
                 command.stdin.write(' ')
 
-        command.stdin.write(';exit\n')
-        if stdin_text:
-            command.stdin.write(stdin_text)
+        command.stdin.write('\nexit\n')
         command.stdin.flush()
         if not bg:
             command.wait(raiseonerr=raiseonerr)
diff --git a/pytest_multihost/transport.py b/pytest_multihost/transport.py
index b8984e3..f45a025 100644
--- a/pytest_multihost/transport.py
+++ b/pytest_multihost/transport.py
@@ -327,7 +327,8 @@ class OpenSSHTransport(Transport):
         command = self._run(['bash'], argv=argv, log_stdout=log_stdout)
         return command
 
-    def _run(self, command, log_stdout=True, argv=None, collect_output=True):
+    def _run(self, command, log_stdout=True, argv=None, collect_output=True,
+             encoding='utf-8'):
         """Run the given command on the remote host
 
         :param command: Command to run (appended to the common SSH invocation)
@@ -341,7 +342,8 @@ class OpenSSHTransport(Transport):
         ssh = SSHCallWrapper(self.ssh_argv + list(command))
         return SSHCommand(ssh, argv, logger_name, log_stdout=log_stdout,
                           collect_output=collect_output,
-                          get_logger=self.host.config.get_logger)
+                          get_logger=self.host.config.get_logger,
+                          encoding=encoding)
 
     def file_exists(self, path):
         self.log.info('STAT %s', path)
@@ -364,7 +366,7 @@ class OpenSSHTransport(Transport):
 
     def get_file_contents(self, filename, encoding=None):
         self.log.info('GET %s', filename)
-        cmd = self._run(['cat', filename], log_stdout=False)
+        cmd = self._run(['cat', filename], log_stdout=False, encoding=None)
         cmd.wait(raiseonerr=False)
         if cmd.returncode == 0:
             result = cmd.stdout_text
@@ -443,8 +445,10 @@ class SSHCommand(Command):
 
         self._ssh.invoke_shell()
 
+        self._is_binary = encoding is None or sys.version_info < (3, 0)
+
         def wrap_file(file, encoding):
-            if encoding is None or sys.version_info < (3, 0):
+            if self._is_binary:
                 return file
             else:
                 return io.TextIOWrapper(file, encoding=encoding)
@@ -463,8 +467,13 @@ class SSHCommand(Command):
         while self.running_threads:
             self.running_threads.pop().join()
 
-        self.stdout_text = ''.join(self._stdout_lines)
-        self.stderr_text = ''.join(self._stderr_lines)
+        if self._is_binary:
+            empty_sep = b''
+        else:
+            empty_sep = u''
+
+        self.stdout_text = empty_sep.join(self._stdout_lines)
+        self.stderr_text = empty_sep.join(self._stderr_lines)
         self.returncode = self._ssh.recv_exit_status()
         self._ssh.close()
 
diff --git a/test_pytestmultihost/test_localhost.py b/test_pytestmultihost/test_localhost.py
index 180aa73..96f9591 100644
--- a/test_pytestmultihost/test_localhost.py
+++ b/test_pytestmultihost/test_localhost.py
@@ -121,6 +121,9 @@ def multihost_badpassword(request, transport_class):
 def _first_command(host):
     """If managed command fails, prints a message to help debugging"""
     try:
+        # Run dummy command first; this should catch spurious SSH messages.
+        host.run_command(['echo', 'hello', 'world'])
+        # Now, run the actual command
         yield
     except (AuthenticationException, CalledProcessError):
         print (
@@ -196,6 +199,40 @@ class TestLocalhost(object):
             host.transport.rmdir(filename)
         assert not os.path.exists(filename)
 
+    def test_escaping(self, multihost, tmpdir):
+        host = multihost.host
+        test_file_path = str(tmpdir.join('testfile.txt'))
+
+        stdin_text = '"test", test, "test", $test, '
+        stdin_text += ''.join(chr(x) for x in range(32, 127))
+        tee = host.run_command(
+            ["tee", test_file_path],
+            stdin_text=stdin_text,
+            raiseonerr=False,
+        )
+        print(tee.stderr_text)
+        assert tee.stdout_text == stdin_text + '\n'
+        with open(test_file_path, "r") as f:
+            assert f.read() == tee.stdout_text
+
+    def test_background(self, multihost, tmpdir):
+        host = multihost.host
+
+        pipe_filename = str(tmpdir.join('test.pipe'))
+
+        with _first_command(host):
+            host.run_command(['mkfifo', pipe_filename])
+
+        cat = host.run_command(['cat', pipe_filename], bg=True)
+        host.run_command('cat > ' + pipe_filename, stdin_text='expected value')
+
+        cat.wait()
+        assert cat.stdout_text == 'expected value\n'
+        assert cat.returncode == 0
+
+
+@pytest.mark.needs_ssh
+class TestLocalhostBadConnection(object):
     def test_reset(self, multihost):
         host = multihost.host
         with _first_command(host):
@@ -224,11 +261,3 @@ class TestLocalhost(object):
         host = multihost_badpassword.host
         with pytest.raises((AuthenticationException, RuntimeError)):
             echo = host.run_command(['echo', 'hello', 'world'])
-
-    def test_background(self, multihost):
-	host = multihost.host
-	run_nc = 'nc -l 12080 > /tmp/filename.out'
-	cmd = host.run_command(run_nc, bg=True, raiseonerr=False)
-	send_file = 'nc localhost 12080 < /root/anaconda-ks.cfg'
-	cmd = host.run_command(send_file)
-	assert cmd.returncode == 0

From ebb7c5a076ea31d1095dfde861948223c458dbee Mon Sep 17 00:00:00 2001
From: Petr Viktorin <pviktori@redhat.com>
Date: Jan 15 2018 15:34:40 +0000
Subject: [PATCH 2/5] Handle encoding/decoding


When writing files or issuing commands, bytestrings are passed through
unchanged, and text strings (`unicode` in Python 2) are encoded using
a configurable encoding (`utf-8` by default).

When reading files, bytestrings are returned by default,
but an encoding can be given to get a test string.

For command output, separate `stdout_bytes` and `stdout_text` attributes
are provided.
The latter uses a configurable encoding (`utf-8` by default).

Fixes: https://pagure.io/python-pytest-multihost/issue/7

---

diff --git a/README.rst b/README.rst
index c33d4d9..09e6d9d 100644
--- a/README.rst
+++ b/README.rst
@@ -172,6 +172,22 @@ $ py.test --multihost-config=/path/to/configfile.yaml
 To use YAML files, the PyYAML package is required. Without it only JSON files
 can be used.
 
+
+Encoding and bytes/text
+-----------------------
+
+When writing files or issuing commands, bytestrings are passed through
+unchanged, and text strings (``unicode`` in Python 2) are encoded using
+a configurable encoding (``utf-8`` by default).
+
+When reading files, bytestrings are returned by default,
+but an encoding can be given to get a test string.
+
+For command output, separate ``stdout_bytes`` and ``stdout_text`` attributes
+are provided.
+The latter uses a configurable encoding (``utf-8` by default).
+
+
 Contributing
 ------------
 
diff --git a/pytest_multihost/host.py b/pytest_multihost/host.py
index ec26496..b1a000d 100644
--- a/pytest_multihost/host.py
+++ b/pytest_multihost/host.py
@@ -25,7 +25,7 @@ class BaseHost(object):
     See README for an overview of the core classes.
     """
     transport_class = transport.SSHTransport
-    command_prelude = ''
+    command_prelude = b''
 
     def __init__(self, domain, hostname, role, ip=None,
                  external_hostname=None, username=None, password=None,
@@ -190,9 +190,9 @@ class BaseHost(object):
         """Shortcut for transport.get_file_contents"""
         return self.transport.get_file_contents(filename, encoding=encoding)
 
-    def put_file_contents(self, filename, contents):
+    def put_file_contents(self, filename, contents, encoding='utf-8'):
         """Shortcut for transport.put_file_contents"""
-        self.transport.put_file_contents(filename, contents)
+        self.transport.put_file_contents(filename, contents, encoding=encoding)
 
     def collect_log(self, filename):
         """Call all registered log collectors on the given filename"""
@@ -201,7 +201,7 @@ class BaseHost(object):
 
     def run_command(self, argv, set_env=True, stdin_text=None,
                     log_stdout=True, raiseonerr=True,
-                    cwd=None, bg=False):
+                    cwd=None, bg=False, encoding='utf-8'):
         """Run the given command on this host
 
         Returns a Command instance. The command will have already run in the
@@ -219,47 +219,67 @@ class BaseHost(object):
                            does not exit with return code 0
         :param cwd: The working directory for the command
         :param bg: If True, runs command in background
+        :param encoding: Encoding for the resulting Command instance's
+                         ``stdout_text`` and ``stderr_text``, and for
+                         ``stdin_text``, ``argv``, etc. if they are not
+                         bytestrings already.
         """
-        command = self.transport.start_shell(argv, log_stdout=log_stdout)
+        def encode(string):
+            if not isinstance(string, bytes):
+                return string.encode(encoding)
+            else:
+                return string
+
+        command = self.transport.start_shell(argv, log_stdout=log_stdout,
+                                             encoding=encoding)
         # Set working directory
         if cwd is None:
             cwd = self.test_dir
-        command.stdin.write('cd %s\n' % shell_quote(cwd))
+        command.stdin.write(b'cd %s\n' % shell_quote(encode(cwd)))
 
         # Set the environment
         if set_env:
-            command.stdin.write('. %s\n' % shell_quote(self.env_sh_path))
+            quoted = shell_quote(encode(self.env_sh_path))
+            command.stdin.write(b'. %s\n' % quoted)
 
         if self.command_prelude:
-            command.stdin.write(self.command_prelude)
+            command.stdin.write(encode(self.command_prelude))
 
         if stdin_text:
-            command.stdin.write("echo -e '")
-            command.stdin.write(stdin_text.replace("'", r"'\''"))
-            command.stdin.write("' | ")
+            command.stdin.write(b"echo -e ")
+            command.stdin.write(_echo_quote(encode(stdin_text)))
+            command.stdin.write(b" | ")
 
         if isinstance(argv, basestring):
             # Run a shell command given as a string
-            command.stdin.write('(')
-            command.stdin.write(argv)
-            command.stdin.write(')')
+            command.stdin.write(b'(')
+            command.stdin.write(encode(argv))
+            command.stdin.write(b')')
         else:
             # Run a command given as a popen-style list (no shell expansion)
             for arg in argv:
-                command.stdin.write(shell_quote(arg))
-                command.stdin.write(' ')
+                command.stdin.write(shell_quote(encode(arg)))
+                command.stdin.write(b' ')
 
-        command.stdin.write('\nexit\n')
+        command.stdin.write(b'\nexit\n')
         command.stdin.flush()
         if not bg:
             command.wait(raiseonerr=raiseonerr)
         return command
 
 
+def _echo_quote(bytestring):
+    """Encode a bytestring for use with bash & "echo -e"
+    """
+    bytestring = bytestring.replace(b"\\", br"\\")
+    bytestring = bytestring.replace(b"\0", br"\x00")
+    bytestring = bytestring.replace(b"'", br"'\''")
+    return b"'" + bytestring + b"'"
+
 
 class Host(BaseHost):
     """A Unix host"""
-    command_prelude = 'set -e\n'
+    command_prelude = b'set -e\n'
 
 
 class WinHost(BaseHost):
diff --git a/pytest_multihost/transport.py b/pytest_multihost/transport.py
index f45a025..f59a8ce 100644
--- a/pytest_multihost/transport.py
+++ b/pytest_multihost/transport.py
@@ -44,11 +44,19 @@ class Transport(object):
         self._command_index = 0
 
     def get_file_contents(self, filename, encoding=None):
-        """Read the named remote file and return the contents as a string"""
+        """Read the named remote file and return the contents
+
+        The string will be decoded using the given encoding;
+        if encoding is None (default), it will be returned as a bytestring.
+        """
         raise NotImplementedError('Transport.get_file_contents')
 
-    def put_file_contents(self, filename, contents):
-        """Write the given string to the named remote file"""
+    def put_file_contents(self, filename, contents, encoding='utf-8'):
+        """Write the given string (or bytestring) to the named remote file
+
+        The contents string will be encoded using the given encoding
+        (default: ``'utf-8'``), unless aleady a bytestring.
+        """
         raise NotImplementedError('Transport.put_file_contents')
 
     def file_exists(self, filename):
@@ -59,18 +67,22 @@ class Transport(object):
         """Make the named directory"""
         raise NotImplementedError('Transport.mkdir')
 
-    def start_shell(self, argv, log_stdout=True):
+    def start_shell(self, argv, log_stdout=True, encoding=None):
         """Start a Shell
 
         :param argv: The command this shell is intended to run (used for
                      logging only)
         :param log_stdout: If false, the stdout will not be logged (useful when
                            binary output is expected)
+        :param encoding: Encoding for the resulting Command's ``stdout_text``
+                         and ``stderr_text``.
 
         Given a `shell` from this method, the caller can then use
         ``shell.stdin.write()`` to input any command(s), call ``shell.wait()``
         to let the command run, and then inspect ``returncode``,
         ``stdout_text`` or ``stderr_text``.
+
+        Note that ``shell.stdin`` uses bytes I/O.
         """
         raise NotImplementedError('Transport.start_shell')
 
@@ -84,7 +96,7 @@ class Transport(object):
 
     def get_file(self, remotepath, localpath):
         """Copy a file from the remote host to a local file"""
-        contents = self.get_file_contents(remotepath)
+        contents = self.get_file_contents(remotepath, encoding=None)
         with open(localpath, 'wb') as local_file:
             local_file.write(contents)
 
@@ -92,7 +104,7 @@ class Transport(object):
         """Copy a local file to the remote host"""
         with open(localpath, 'rb') as local_file:
             contents = local_file.read()
-        self.put_file_contents(remotepath, contents)
+        self.put_file_contents(remotepath, contents, encoding=None)
 
     def get_next_command_logger_name(self):
         self._command_index += 1
@@ -111,6 +123,28 @@ class Transport(object):
         raise NotImplementedError('Transport.remove_file')
 
 
+class _decoded_output_property(object):
+    """Descriptor for on-demand decoding of a Command's output stream
+    """
+    def __init__(self, name):
+        self.name = name
+
+    def __set_name__(self, cls, name):
+        # Sanity check (called only on Python 3.6+).
+        # This property expects to handle attributes named '<foo>_text'.
+        assert name == self.name + '_text'
+
+    def __get__(self, instance, cls=None):
+        if instance is None:
+            return self
+        else:
+            bytestring = getattr(instance, self.name + '_bytes')
+            print(bytestring, instance.encoding)
+            decoded = bytestring.decode(instance.encoding)
+            setattr(instance, self.name + '_text', decoded)
+            return decoded
+
+
 class Command(object):
     """A Popen-style object representing a remote command
 
@@ -122,12 +156,17 @@ class Command(object):
     To make sure reading doesn't stall after one buffer fills up, they are read
     in parallel using threads.
 
-    After calling wait(), ``stdout_text`` and ``stderr_text`` attributes will
-    be strings containing the output, and ``returncode`` will contain the
+    After calling wait(), ``stdout_bytes`` and ``stderr_bytes`` attributes will
+    be bytestrings containing the output, and ``returncode`` will contain the
     exit code.
+
+    The ``stdout_text`` and ``stdout_text`` will be the corresponding output
+    decoded using the given ``encoding`` (default: ``'utf-8'``).
+    These are decoded on-demand; do not access them if a command
+    produces binary output.
     """
     def __init__(self, argv, logger_name=None, log_stdout=True,
-                 get_logger=None):
+                 get_logger=None, encoding='utf-8'):
         self.returncode = None
         self.argv = argv
         self._done = False
@@ -140,6 +179,10 @@ class Command(object):
             get_logger = logging.getLogger
         self.get_logger = get_logger
         self.log = get_logger(self.logger_name)
+        self.encoding = encoding
+
+    stdout_text = _decoded_output_property('stdout')
+    stderr_text = _decoded_output_property('stderr')
 
     def wait(self, raiseonerr=True):
         """Wait for the remote process to exit
@@ -220,16 +263,18 @@ class ParamikoTransport(Transport):
     def get_file_contents(self, filename, encoding=None):
         """Read the named remote file and return the contents as a string"""
         self.log.debug('READ %s', filename)
-        with self.sftp_open(filename) as f:
+        with self.sftp_open(filename, 'rb') as f:
             result = f.read()
         if encoding:
             result = result.decode(encoding)
         return result
 
-    def put_file_contents(self, filename, contents):
+    def put_file_contents(self, filename, contents, encoding=None):
         """Write the given string to the named remote file"""
         self.log.info('WRITE %s', filename)
-        with self.sftp_open(filename, 'w') as f:
+        if encoding and not isinstance(contents, bytes):
+            contents = contents.encode(encoding)
+        with self.sftp_open(filename, 'wb') as f:
             f.write(contents)
 
     def file_exists(self, filename):
@@ -248,13 +293,14 @@ class ParamikoTransport(Transport):
         self.log.info('MKDIR %s', path)
         self.sftp.mkdir(path)
 
-    def start_shell(self, argv, log_stdout=True):
+    def start_shell(self, argv, log_stdout=True, encoding='utf-8'):
         logger_name = self.get_next_command_logger_name()
         ssh = self._transport.open_channel('session')
         self.log.info('RUN %s', argv)
         return SSHCommand(ssh, argv, logger_name=logger_name,
                           log_stdout=log_stdout,
-                          get_logger=self.host.config.get_logger)
+                          get_logger=self.host.config.get_logger,
+                          encoding=encoding)
 
     def get_file(self, remotepath, localpath):
         self.log.debug('GET %s', remotepath)
@@ -322,9 +368,10 @@ class OpenSSHTransport(Transport):
 
         return argv
 
-    def start_shell(self, argv, log_stdout=True):
+    def start_shell(self, argv, log_stdout=True, encoding='utf-8'):
         self.log.info('RUN %s', argv)
-        command = self._run(['bash'], argv=argv, log_stdout=log_stdout)
+        command = self._run(['bash'], argv=argv, log_stdout=log_stdout,
+                            encoding=encoding)
         return command
 
     def _run(self, command, log_stdout=True, argv=None, collect_output=True,
@@ -357,19 +404,21 @@ class OpenSSHTransport(Transport):
         cmd = self._run(['mkdir', path])
         cmd.wait()
 
-    def put_file_contents(self, filename, contents):
+    def put_file_contents(self, filename, contents, encoding='utf-8'):
         self.log.info('PUT %s', filename)
+        if encoding and not isinstance(contents, bytes):
+            contents = contents.encode(encoding)
         cmd = self._run(['tee', filename], log_stdout=False)
         cmd.stdin.write(contents)
         cmd.wait()
-        assert cmd.stdout_text == contents
+        assert cmd.stdout_bytes == contents
 
     def get_file_contents(self, filename, encoding=None):
         self.log.info('GET %s', filename)
-        cmd = self._run(['cat', filename], log_stdout=False, encoding=None)
+        cmd = self._run(['cat', filename], log_stdout=False)
         cmd.wait(raiseonerr=False)
         if cmd.returncode == 0:
-            result = cmd.stdout_text
+            result = cmd.stdout_bytes
             if encoding:
                 result = result.decode(encoding)
             return result
@@ -434,7 +483,8 @@ class SSHCommand(Command):
                  collect_output=True, encoding='utf-8', get_logger=None):
         super(SSHCommand, self).__init__(argv, logger_name,
                                          log_stdout=log_stdout,
-                                         get_logger=get_logger)
+                                         get_logger=get_logger,
+                                         encoding=encoding)
         self._stdout_lines = []
         self._stderr_lines = []
         self.running_threads = set()
@@ -445,16 +495,16 @@ class SSHCommand(Command):
 
         self._ssh.invoke_shell()
 
-        self._is_binary = encoding is None or sys.version_info < (3, 0)
+        self._use_bytes = (encoding is None)
 
         def wrap_file(file, encoding):
-            if self._is_binary:
+            if self._use_bytes:
                 return file
             else:
                 return io.TextIOWrapper(file, encoding=encoding)
-        self.stdin = wrap_file(self._ssh.makefile('wb'), 'utf-8')
-        stdout = wrap_file(self._ssh.makefile('rb'), encoding)
-        stderr = wrap_file(self._ssh.makefile_stderr('rb'), encoding)
+        self.stdin = self._ssh.makefile('wb')
+        stdout = self._ssh.makefile('rb')
+        stderr = self._ssh.makefile_stderr('rb')
 
         if collect_output:
             self._start_pipe_thread(self._stdout_lines, stdout, 'out',
@@ -467,13 +517,9 @@ class SSHCommand(Command):
         while self.running_threads:
             self.running_threads.pop().join()
 
-        if self._is_binary:
-            empty_sep = b''
-        else:
-            empty_sep = u''
+        self.stdout_bytes = b''.join(self._stdout_lines)
+        self.stderr_bytes = b''.join(self._stderr_lines)
 
-        self.stdout_text = empty_sep.join(self._stdout_lines)
-        self.stderr_text = empty_sep.join(self._stderr_lines)
         self.returncode = self._ssh.recv_exit_status()
         self._ssh.close()
 
@@ -489,7 +535,8 @@ class SSHCommand(Command):
         def read_stream():
             for line in stream:
                 if do_log:
-                    log.debug(line.rstrip('\n'))
+                    log.debug(line.rstrip(b'\n').decode('utf-8',
+                                                        errors='replace'))
                 result_list.append(line)
 
         thread = threading.Thread(target=read_stream)
diff --git a/pytest_multihost/util.py b/pytest_multihost/util.py
index 531e999..46889f9 100644
--- a/pytest_multihost/util.py
+++ b/pytest_multihost/util.py
@@ -15,9 +15,9 @@ def check_config_dict_empty(dct, name):
                          (name, ', '.join(dct)))
 
 
-def shell_quote(string):
-    """Quotes a string for the Bash shell"""
-    return "'" + string.replace("'", "'\\''") + "'"
+def shell_quote(bytestring):
+    """Quotes a bytestring for the Bash shell"""
+    return b"'" + bytestring.replace(b"'", b"'\\''") + b"'"
 
 
 class TempDir(object):
diff --git a/test_pytestmultihost/test_localhost.py b/test_pytestmultihost/test_localhost.py
index 96f9591..134c1c2 100644
--- a/test_pytestmultihost/test_localhost.py
+++ b/test_pytestmultihost/test_localhost.py
@@ -6,6 +6,7 @@ import getpass
 import pytest
 from subprocess import CalledProcessError
 import contextlib
+import sys
 import os
 
 import pytest_multihost
@@ -162,6 +163,38 @@ class TestLocalhost(object):
         with pytest.raises(IOError):
             host.get_file_contents(filename)
 
+    def test_get_put_file_contents_bytes(self, multihost, tmpdir):
+        host = multihost.host
+        filename = str(tmpdir.join('test-bytes.txt'))
+        testbytes = u'test \0 \N{WHITE SMILING FACE}'.encode('utf-8')
+        with _first_command(host):
+            host.put_file_contents(filename, testbytes, encoding=None)
+        result = host.get_file_contents(filename, encoding=None)
+        assert result == testbytes
+
+    @pytest.mark.parametrize('encoding', ('utf-8', 'utf-16'))
+    def test_put_file_contents_utf(self, multihost, tmpdir, encoding):
+        host = multihost.host
+        filename = str(tmpdir.join('test-{}.txt'.format(encoding)))
+        teststring = u'test \N{WHITE SMILING FACE}'
+        with _first_command(host):
+            host.put_file_contents(filename, teststring, encoding=encoding)
+        result = host.get_file_contents(filename, encoding=None)
+        assert result == teststring.encode(encoding)
+        with open(filename, 'rb') as f:
+            assert f.read() == teststring.encode(encoding)
+
+    @pytest.mark.parametrize('encoding', ('utf-8', 'utf-16'))
+    def test_get_file_contents_encoding(self, multihost, tmpdir, encoding):
+        host = multihost.host
+        filename = str(tmpdir.join('test-{}.txt'.format(encoding)))
+        teststring = u'test \N{WHITE SMILING FACE}'
+        with open(filename, 'wb') as f:
+            f.write(teststring.encode(encoding))
+        result = host.get_file_contents(filename, encoding=encoding)
+        assert result == teststring
+        assert type(result) == type(u'')
+
     def test_rename_file(self, multihost, tmpdir):
         host = multihost.host
         filename = str(tmpdir.join('test.txt'))
@@ -205,6 +238,8 @@ class TestLocalhost(object):
 
         stdin_text = '"test", test, "test", $test, '
         stdin_text += ''.join(chr(x) for x in range(32, 127))
+        stdin_text += r', \x66\0111\x00, '
+        stdin_text += ''.join('\\' + chr(x) for x in range(32, 127))
         tee = host.run_command(
             ["tee", test_file_path],
             stdin_text=stdin_text,
@@ -215,6 +250,22 @@ class TestLocalhost(object):
         with open(test_file_path, "r") as f:
             assert f.read() == tee.stdout_text
 
+    def test_escaping_binary(self, multihost, tmpdir):
+        host = multihost.host
+        test_file_path = str(tmpdir.join('testfile.txt'))
+
+        stdin_bytes = b'"test", test, "test", $test, '
+        stdin_bytes += bytes(range(0, 256))
+        stdin_bytes += br', \x66\0111\x00'
+        tee = host.run_command(
+            ["tee", test_file_path],
+            stdin_text=stdin_bytes,
+            raiseonerr=False,
+        )
+        assert tee.stdout_bytes == stdin_bytes + b'\n'
+        with open(test_file_path, "rb") as f:
+            assert f.read() == tee.stdout_bytes
+
     def test_background(self, multihost, tmpdir):
         host = multihost.host
 

From 1d52603faa328bdd5407aebfa52c8b9b63617eb0 Mon Sep 17 00:00:00 2001
From: Petr Viktorin <pviktori@redhat.com>
Date: Jan 15 2018 15:34:40 +0000
Subject: [PATCH 3/5] Make Command a context manager; document commands must be waited for


Fixes: https://pagure.io/python-pytest-multihost/issue/6

---

diff --git a/pytest_multihost/host.py b/pytest_multihost/host.py
index b1a000d..3a488f4 100644
--- a/pytest_multihost/host.py
+++ b/pytest_multihost/host.py
@@ -218,7 +218,10 @@ class BaseHost(object):
         :param raiseonerr: If true, an exception will be raised if the command
                            does not exit with return code 0
         :param cwd: The working directory for the command
-        :param bg: If True, runs command in background
+        :param bg: If True, runs command in background.
+                   In this case, either the result should be used in a ``with``
+                   statement, or ``wait()`` should be called explicitly
+                   when the command is finished.
         :param encoding: Encoding for the resulting Command instance's
                          ``stdout_text`` and ``stderr_text``, and for
                          ``stdin_text``, ``argv``, etc. if they are not
@@ -263,8 +266,9 @@ class BaseHost(object):
 
         command.stdin.write(b'\nexit\n')
         command.stdin.flush()
+        command.raiseonerr = raiseonerr
         if not bg:
-            command.wait(raiseonerr=raiseonerr)
+            command.wait()
         return command
 
 
diff --git a/pytest_multihost/transport.py b/pytest_multihost/transport.py
index f59a8ce..2d647d2 100644
--- a/pytest_multihost/transport.py
+++ b/pytest_multihost/transport.py
@@ -164,6 +164,11 @@ class Command(object):
     decoded using the given ``encoding`` (default: ``'utf-8'``).
     These are decoded on-demand; do not access them if a command
     produces binary output.
+
+    A Command may be used as a context manager (in the ``with`` statement).
+    Exiting the context will automatically call ``wait()``.
+    This raises an exception if the exit code is not 0, unless the
+    ``raiseonerr`` attribute is set to false before exiting the context.
     """
     def __init__(self, argv, logger_name=None, log_stdout=True,
                  get_logger=None, encoding='utf-8'):
@@ -180,6 +185,7 @@ class Command(object):
         self.get_logger = get_logger
         self.log = get_logger(self.logger_name)
         self.encoding = encoding
+        self.raiseonerr = True
 
     stdout_text = _decoded_output_property('stdout')
     stderr_text = _decoded_output_property('stderr')
@@ -187,7 +193,7 @@ class Command(object):
     def wait(self, raiseonerr=True):
         """Wait for the remote process to exit
 
-        Raises an excption if the exit code is not 0, unless raiseonerr is
+        Raises an exception if the exit code is not 0, unless raiseonerr is
         true.
         """
         if self._done:
@@ -211,6 +217,12 @@ class Command(object):
         """
         raise NotImplementedError()
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc_info):
+        self.wait(raiseonerr=self.raiseonerr)
+
 
 class ParamikoTransport(Transport):
     """Transport that uses the Paramiko SSH2 library"""
diff --git a/test_pytestmultihost/test_localhost.py b/test_pytestmultihost/test_localhost.py
index 134c1c2..5744997 100644
--- a/test_pytestmultihost/test_localhost.py
+++ b/test_pytestmultihost/test_localhost.py
@@ -266,7 +266,7 @@ class TestLocalhost(object):
         with open(test_file_path, "rb") as f:
             assert f.read() == tee.stdout_bytes
 
-    def test_background(self, multihost, tmpdir):
+    def test_background_explicit_wait(self, multihost, tmpdir):
         host = multihost.host
 
         pipe_filename = str(tmpdir.join('test.pipe'))
@@ -281,6 +281,21 @@ class TestLocalhost(object):
         assert cat.stdout_text == 'expected value\n'
         assert cat.returncode == 0
 
+    def test_background_context(self, multihost, tmpdir):
+        host = multihost.host
+
+        pipe_filename = str(tmpdir.join('test.pipe'))
+
+        with _first_command(host):
+            host.run_command(['mkfifo', pipe_filename])
+
+        with host.run_command(['cat', pipe_filename], bg=True) as cat:
+            host.run_command('cat > ' + pipe_filename,
+                             stdin_text='expected value')
+
+        assert cat.stdout_text == 'expected value\n'
+        assert cat.returncode == 0
+
 
 @pytest.mark.needs_ssh
 class TestLocalhostBadConnection(object):

From 680d0ca41a7c9ac4c279b6d81f9588cd305ce5d0 Mon Sep 17 00:00:00 2001
From: Petr Viktorin <pviktori@redhat.com>
Date: Jan 15 2018 15:34:40 +0000
Subject: [PATCH 4/5] Add a tox.ini to ease testing on various Python versions


---

diff --git a/tox.ini b/tox.ini
new file mode 100644
index 0000000..f57359a
--- /dev/null
+++ b/tox.ini
@@ -0,0 +1,14 @@
+# Tox (http://tox.testrun.org/) is a tool for running tests
+# in multiple virtualenvs. This configuration file will run the
+# test suite on all supported python versions. To use it, "pip install tox"
+# and then run "tox" from this directory.
+
+[tox]
+envlist = py2,py36
+minver = 1.8
+
+[testenv]
+deps =
+    pytest
+    paramiko
+commands = python -m pytest -vv test_pytestmultihost/

From 6994d275ec5468a73c8a5c4670e6003c8d3227d1 Mon Sep 17 00:00:00 2001
From: Petr Viktorin <pviktori@redhat.com>
Date: Feb 12 2018 09:27:28 +0000
Subject: [PATCH 5/5] Honor instance attribute if `raiseonerr` is not passed to Command.wait()


---

diff --git a/pytest_multihost/transport.py b/pytest_multihost/transport.py
index 2d647d2..e1bac45 100644
--- a/pytest_multihost/transport.py
+++ b/pytest_multihost/transport.py
@@ -30,6 +30,8 @@ except ImportError:
     have_paramiko = False
 
 
+DEFAULT = object()
+
 class Transport(object):
     """Mechanism for communicating with remote hosts
 
@@ -190,12 +192,18 @@ class Command(object):
     stdout_text = _decoded_output_property('stdout')
     stderr_text = _decoded_output_property('stderr')
 
-    def wait(self, raiseonerr=True):
+    def wait(self, raiseonerr=DEFAULT):
         """Wait for the remote process to exit
 
-        Raises an exception if the exit code is not 0, unless raiseonerr is
+        Raises an exception if the exit code is not 0, unless ``raiseonerr`` is
         true.
+
+        When ``raiseonerr`` is not specified as argument, the ``raiseonerr``
+        attribute is used.
         """
+        if raiseonerr is DEFAULT:
+            raiseonerr = self.raiseonerr
+
         if self._done:
             return self.returncode
 
diff --git a/test_pytestmultihost/test_localhost.py b/test_pytestmultihost/test_localhost.py
index 5744997..22f7e62 100644
--- a/test_pytestmultihost/test_localhost.py
+++ b/test_pytestmultihost/test_localhost.py
@@ -297,6 +297,31 @@ class TestLocalhost(object):
         assert cat.returncode == 0
 
 
+    def test_background_raiseonerr_false(self, multihost, tmpdir):
+        host = multihost.host
+        with _first_command(host):
+            false = host.run_command(['false'], raiseonerr=False, bg=True)
+
+        assert false.returncode != 0
+
+
+    def test_background_raiseonerr_with(self, multihost, tmpdir):
+        host = multihost.host
+        with _first_command(host):
+            with pytest.raises(CalledProcessError):
+                with host.run_command(['false'], raiseonerr=True, bg=True):
+                    pass
+
+    def test_background_raiseonerr_wait(self, multihost, tmpdir):
+        host = multihost.host
+        with _first_command(host):
+            false = host.run_command(['false'], raiseonerr=True, bg=True)
+
+            with pytest.raises(CalledProcessError):
+                false.wait()
+
+
+
 @pytest.mark.needs_ssh
 class TestLocalhostBadConnection(object):
     def test_reset(self, multihost):