Merge lp:~frankban/python-shelltoolbox/helpers into lp:python-shelltoolbox

[Francesco Banconi]

== Changes ==

- Added several helper functions:
  - bzr_whois
  - file_append
  - file_prepend
  - generate_ssh_keys
  - get_su_command
  - join_command
  - mkdirs
  - ssh
  - user_exists

- Added `environ` context manager and updated `su` to use it.

- Changed `apt_get_install` helper: now the function uses environ to run dpkg
  in non-interactive mode.

- Changed `get_user_home` to handle non-existent users
  (returning a default home directory)

- Updated the `install_extra_repositories` helper to accept distribution
  placeholders.

- Updated `cd` context manager: yield in a try/finally block.

- Fixed `grep` helper.

== Tests ==

python tests.py
.....................................................
----------------------------------------------------------------------
Ran 53 tests in 0.124s

OK

[Gary Poster (gary)]

Thank you! Small comments and suggestions follow.

I asked about reordering the functions for alphabetization. Not having run at the top is somewhat surprising to me, but I am ok with it.

I was struck that file_append really expects the line to be added to end with a \n. I suppose we could enforce that one way or another...or not. I'm OK with leaving as is, or you could add a \n if the line does not already end with one (and do this before you check if the line is in the content). On a somewhat related note, it might be better to check if the line to be added is in the output of readlines(), not read(); that would verify that the line, not a line that ends with the line, is in the file. For example, right now a file of "abc\ndef\n" would incorrectly show that the "ef\n" line was in the file.

In file_prepend, I am tempted to suggest that this snippet should be changed:
            if line in lines:
                lines.remove(line)
This iterates through the lines twice, which is mildly annoying to me but almost certainly not a practical concern for us with these use cases. That said, I'd be tempted to change that to the following:
            try:
                lines.remove(line)
            except ValueError:
                pass
You could also verify that the line to be added in this function ended with a \n if you wanted to.

I don't really understand this part of the docstring for get_su_command:
    This can be used together with `run` when the `su` context manager is not
    enough (e.g. an external program uses uid rather than euid).
Oh! You mean that you would use run(*get_su_command(...)). Maybe an example would help.

I'm a bit concerned about what you are doing here in get_value_from_line:
    return line.split('=')[1].strip('"\' ')
would maybe json be a safer approach? No, I tried it, and json.loads('"hi"') is fine but json.loads("'hi'") is not. :-/ I dunno. When do we use this? The logic seems a bit idiosyncratic.

The grep command also seems a bit idiosyncratic. Why is it reasonable to always strip the result? I also would be tempted to rename the funtion: grep does a lot more than that. simply search_file? Hm, I see that you just reordered this file, you didn't actually add grep.... Never mind then, I guess. I should have commented on this sooner, but you shouldn't have to worry about it. If you *do* want to change it, please leave the "grep" alias around for now so that our charms do not suddenly break.

I don't understand why you removed the Serializer. Was it defined twice?

I thought the choices you made for doc examples and for unit tests were good. Thank you.

Great job!

Gary

review: Approve

[Francesco Banconi (frankban)]

> I was struck that file_append really expects the line to be added to end with
> a \n. I suppose we could enforce that one way or another...or not. I'm OK
> with leaving as is, or you could add a \n if the line does not already end
> with one (and do this before you check if the line is in the content). On a
> somewhat related note, it might be better to check if the line to be added is
> in the output of readlines(), not read(); that would verify that the line, not
> a line that ends with the line, is in the file. For example, right now a file
> of "abc\ndef\n" would incorrectly show that the "ef\n" line was in the file.

I agree with your suggestions. I've changed the function and removed those
ambiguities. I've added tests for line not ending with '\n' and for line
fragments.

> In file_prepend, I am tempted to suggest that this snippet should be changed:
> if line in lines:
> lines.remove(line)
> This iterates through the lines twice, which is mildly annoying to me but
> almost certainly not a practical concern for us with these use cases. That
> said, I'd be tempted to change that to the following:
> try:
> lines.remove(line)
> except ValueError:
> pass
> You could also verify that the line to be added in this function ended with a
> \n if you wanted to.

Done.

> I'm a bit concerned about what you are doing here in get_value_from_line:
> return line.split('=')[1].strip('"\' ')
> would maybe json be a safer approach? No, I tried it, and json.loads('"hi"')
> is fine but json.loads("'hi'") is not. :-/ I dunno. When do we use this?
> The logic seems a bit idiosyncratic.

I don't remember the origin of that function. Since it is not used in the
charms or in setuplxc/lpsetup, I am going to remove it.

> The grep command also seems a bit idiosyncratic. Why is it reasonable to
> always strip the result? I also would be tempted to rename the funtion: grep
> does a lot more than that. simply search_file? Hm, I see that you just
> reordered this file, you didn't actually add grep.... Never mind then, I
> guess. I should have commented on this sooner, but you shouldn't have to
> worry about it. If you *do* want to change it, please leave the "grep" alias
> around for now so that our charms do not suddenly break.

Actually that helper is not used, and I remember it was before we changed the
way the buildslave is reconfigured. I think it's safe to rename it and change
the returned value.

> I don't understand why you removed the Serializer. Was it defined twice?

Yes it was.

Thanks Gary.

[lp:~frankban/python-shelltoolbox/helpers]

33. By Francesco Banconi on 2012-03-07

Fixed file_append.

34. By Francesco Banconi on 2012-03-07

Fixed file_prepend.

35. By Francesco Banconi on 2012-03-07

Removed get_value_from_line.

36. By Francesco Banconi on 2012-03-07

s/grep/search_file/

=== modified file 'setup.py'
--- setup.py	2012-03-01 20:49:48 +0000
+++ setup.py	2012-03-07 10:19:06 +0000
@@ -9,7 +9,7 @@
 ez_setup.use_setuptools()
 
 
-__version__ = '0.1.1'
+__version__ = '0.2.0'
 
 from setuptools import setup
 

=== modified file 'shelltoolbox/__init__.py'
--- shelltoolbox/__init__.py	2012-03-02 14:59:46 +0000
+++ shelltoolbox/__init__.py	2012-03-07 10:19:06 +0000
@@ -19,25 +19,38 @@
 __metaclass__ = type
 __all__ = [
     'apt_get_install',
+    'bzr_whois',
     'cd',
     'command',
     'DictDiffer',
+    'environ',
+    'file_append',
+    'file_prepend',
+    'generate_ssh_keys',
+    'get_su_command',
+    'get_user_home',
     'get_user_ids',
-    'get_user_home',
-    'get_value_from_line',
-    'grep',
     'install_extra_repositories',
+    'join_command',
+    'mkdirs',
     'run',
     'Serializer',
     'script_name',
+    'search_file',
+    'ssh',
     'su',
+    'user_exists',
+    'wait_for_page_contents',
     ]
 
 from collections import namedtuple
 from contextlib import contextmanager
+from email.Utils import parseaddr
+import errno
 import json
 import operator
 import os
+import pipes
 import pwd
 import re
 import subprocess
@@ -50,26 +63,50 @@
 Env = namedtuple('Env', 'uid gid home')
 
 
-def run(*args, **kwargs):
-    """Run the command with the given arguments.
-
-    The first argument is the path to the command to run.
-    Subsequent arguments are command-line arguments to be passed.
-
-    This function accepts all optional keyword arguments accepted by
-    `subprocess.Popen`.
-    """
-    args = [i for i in args if i is not None]
-    pipe = subprocess.PIPE
-    process = subprocess.Popen(
-        args, stdout=kwargs.pop('stdout', pipe),
-        stderr=kwargs.pop('stderr', pipe),
-        close_fds=kwargs.pop('close_fds', True), **kwargs)
-    stdout, stderr = process.communicate()
-    if process.returncode:
-        raise subprocess.CalledProcessError(
-            process.returncode, repr(args), output=stdout+stderr)
-    return stdout
+def apt_get_install(*args, **kwargs):
+    """Install given packages using apt.
+
+    It is possible to pass environment variables to be set during install
+    using keyword arguments.
+
+    :raises: subprocess.CalledProcessError
+    """
+    debian_frontend = kwargs.pop('DEBIAN_FRONTEND', 'noninteractive')
+    with environ(DEBIAN_FRONTEND=debian_frontend, **kwargs):
+        cmd = ('apt-get', '-y', 'install') + args
+        return run(*cmd)
+
+
+def bzr_whois(user):
+    """Return full name and email of bzr `user`.
+
+    Return None if the given `user` does not have a bzr user id.
+    """
+    with su(user):
+        try:
+            whoami = run('bzr', 'whoami')
+        except (subprocess.CalledProcessError, OSError):
+            return None
+    return parseaddr(whoami)
+
+
+@contextmanager
+def cd(directory):
+    """A context manager to temporarily change current working dir, e.g.::
+
+        >>> import os
+        >>> os.chdir('/tmp')
+        >>> with cd('/bin'): print os.getcwd()
+        /bin
+        >>> print os.getcwd()
+        /tmp
+    """
+    cwd = os.getcwd()
+    os.chdir(directory)
+    try:
+        yield
+    finally:
+        os.chdir(cwd)
 
 
 def command(*base_args):
@@ -97,37 +134,393 @@
     return callable_command
 
 
-apt_get_install = command('apt-get', 'install', '-y', '--force-yes')
+@contextmanager
+def environ(**kwargs):
+    """A context manager to temporarily change environment variables.
+
+    If an existing environment variable is changed, it is restored during
+    context cleanup::
+
+        >>> import os
+        >>> os.environ['MY_VARIABLE'] = 'foo'
+        >>> with environ(MY_VARIABLE='bar'): print os.getenv('MY_VARIABLE')
+        bar
+        >>> print os.getenv('MY_VARIABLE')
+        foo
+        >>> del os.environ['MY_VARIABLE']
+
+    If we are adding environment variables, they are removed during context
+    cleanup::
+
+        >>> import os
+        >>> with environ(MY_VAR1='foo', MY_VAR2='bar'):
+        ...     print os.getenv('MY_VAR1'), os.getenv('MY_VAR2')
+        foo bar
+        >>> os.getenv('MY_VAR1') == os.getenv('MY_VAR2') == None
+        True
+    """
+    backup = {}
+    for key, value in kwargs.items():
+        backup[key] = os.getenv(key)
+        os.environ[key] = value
+    try:
+        yield
+    finally:
+        for key, value in backup.items():
+            if value is None:
+                del os.environ[key]
+            else:
+                os.environ[key] = value
+
+
+def file_append(filename, line):
+    r"""Append given `line`, if not present, at the end of `filename`.
+
+    Usage example::
+
+        >>> import tempfile
+        >>> f = tempfile.NamedTemporaryFile('w', delete=False)
+        >>> f.write('line1\n')
+        >>> f.close()
+        >>> file_append(f.name, 'new line\n')
+        >>> open(f.name).read()
+        'line1\nnew line\n'
+
+    Nothing happens if the file already contains the given `line`::
+
+        >>> file_append(f.name, 'new line\n')
+        >>> open(f.name).read()
+        'line1\nnew line\n'
+
+    A new line is automatically added before the given `line` if it is not
+    present at the end of current file content::
+
+        >>> import tempfile
+        >>> f = tempfile.NamedTemporaryFile('w', delete=False)
+        >>> f.write('line1')
+        >>> f.close()
+        >>> file_append(f.name, 'new line\n')
+        >>> open(f.name).read()
+        'line1\nnew line\n'
+
+    The file is created if it does not exist::
+
+        >>> import tempfile
+        >>> filename = tempfile.mktemp()
+        >>> file_append(filename, 'line1\n')
+        >>> open(filename).read()
+        'line1\n'
+    """
+    if not line.endswith('\n'):
+        line += '\n'
+    with open(filename, 'a+') as f:
+        lines = f.readlines()
+        if line not in lines:
+            if not lines or lines[-1].endswith('\n'):
+                f.write(line)
+            else:
+                f.write('\n' + line)
+
+
+def file_prepend(filename, line):
+    r"""Insert given `line`, if not present, at the beginning of `filename`.
+
+    Usage example::
+
+        >>> import tempfile
+        >>> f = tempfile.NamedTemporaryFile('w', delete=False)
+        >>> f.write('line1\n')
+        >>> f.close()
+        >>> file_prepend(f.name, 'line0\n')
+        >>> open(f.name).read()
+        'line0\nline1\n'
+
+    If the file starts with the given `line`, nothing happens::
+
+        >>> file_prepend(f.name, 'line0\n')
+        >>> open(f.name).read()
+        'line0\nline1\n'
+
+    If the file contains the given `line`, but not at the beginning,
+    the line is moved on top::
+
+        >>> file_prepend(f.name, 'line1\n')
+        >>> open(f.name).read()
+        'line1\nline0\n'
+    """
+    if not line.endswith('\n'):
+        line += '\n'
+    with open(filename, 'r+') as f:
+        lines = f.readlines()
+        if lines[0] != line:
+            try:
+                lines.remove(line)
+            except ValueError:
+                pass
+            lines.insert(0, line)
+            f.seek(0)
+            f.writelines(lines)
+
+
+def generate_ssh_keys(path, passphrase=''):
+    """Generate ssh key pair, saving them inside the given `directory`.
+
+        >>> generate_ssh_keys('/tmp/id_rsa')
+        0
+        >>> open('/tmp/id_rsa').readlines()[0].strip()
+        '-----BEGIN RSA PRIVATE KEY-----'
+        >>> open('/tmp/id_rsa.pub').read().startswith('ssh-rsa')
+        True
+        >>> os.remove('/tmp/id_rsa')
+        >>> os.remove('/tmp/id_rsa.pub')
+    """
+    return subprocess.call([
+        'ssh-keygen', '-q', '-t', 'rsa', '-N', passphrase, '-f', path])
+
+
+def get_su_command(user, args):
+    """Return a command line as a sequence, prepending "su" if necessary.
+
+    This can be used together with `run` when the `su` context manager is not
+    enough (e.g. an external program uses uid rather than euid).
+
+        run(*get_su_command(user, ['bzr', 'whoami']))
+
+    If the su is requested as current user, the arguments are returned as
+    given::
+
+        >>> import getpass
+        >>> current_user = getpass.getuser()
+
+        >>> get_su_command(current_user, ('ls', '-l'))
+        ('ls', '-l')
+
+    Otherwise, "su" is prepended::
+
+        >>> get_su_command('nobody', ('ls', '-l', 'my file'))
+        ('su', 'nobody', '-c', "ls -l 'my file'")
+    """
+    if get_user_ids(user)[0] != os.getuid():
+        args = [i for i in args if i is not None]
+        return ('su', user, '-c', join_command(args))
+    return args
+
+
+def get_user_home(user):
+    """Return the home directory of the given `user`.
+
+        >>> get_user_home('root')
+        '/root'
+
+    If the user does not exist, return a default /home/[username] home::
+
+        >>> get_user_home('_this_user_does_not_exist_')
+        '/home/_this_user_does_not_exist_'
+    """
+    try:
+        return pwd.getpwnam(user).pw_dir
+    except KeyError:
+        return os.path.join(os.path.sep, 'home', user)
+
+
+def get_user_ids(user):
+    """Return the uid and gid of given `user`, e.g.::
+
+        >>> get_user_ids('root')
+        (0, 0)
+    """
+    userdata = pwd.getpwnam(user)
+    return userdata.pw_uid, userdata.pw_gid
 
 
 def install_extra_repositories(*repositories):
     """Install all of the extra repositories and update apt.
 
+    Given repositories can contain a "{distribution}" placeholder, that will
+    be replaced by current distribution codename.
+
     :raises: subprocess.CalledProcessError
     """
     distribution = run('lsb_release', '-cs').strip()
-    # Starting from Oneiric, the `apt-add-repository` is interactive by
+    # Starting from Oneiric, `apt-add-repository` is interactive by
     # default, and requires a "-y" flag to be set.
     assume_yes = None if distribution == 'lucid' else '-y'
     for repo in repositories:
-        run('apt-add-repository', assume_yes, repo)
+        repository = repo.format(distribution=distribution)
+        run('apt-add-repository', assume_yes, repository)
     run('apt-get', 'clean')
     run('apt-get', 'update')
 
 
-def grep(content, filename):
+def join_command(args):
+    """Return a valid Unix command line from `args`.
+
+        >>> join_command(['ls', '-l'])
+        'ls -l'
+
+    Arguments containing spaces and empty args are correctly quoted::
+
+        >>> join_command(['command', 'arg1', 'arg containing spaces', ''])
+        "command arg1 'arg containing spaces' ''"
+    """
+    return ' '.join(pipes.quote(arg) for arg in args)
+
+
+def mkdirs(*args):
+    """Create leaf directories (given as `args`) and all intermediate ones.
+
+        >>> import tempfile
+        >>> base_dir = tempfile.mktemp(suffix='/')
+        >>> dir1 = tempfile.mktemp(prefix=base_dir)
+        >>> dir2 = tempfile.mktemp(prefix=base_dir)
+        >>> mkdirs(dir1, dir2)
+        >>> os.path.isdir(dir1)
+        True
+        >>> os.path.isdir(dir2)
+        True
+
+    If the leaf directory already exists the function returns without errors::
+
+        >>> mkdirs(dir1)
+
+    An `OSError` is raised if the leaf path exists and it is a file::
+
+        >>> f = tempfile.NamedTemporaryFile(
+        ...     'w', delete=False, prefix=base_dir)
+        >>> f.close()
+        >>> mkdirs(f.name) # doctest: +ELLIPSIS
+        Traceback (most recent call last):
+        OSError: ...
+    """
+    for directory in args:
+        try:
+            os.makedirs(directory)
+        except OSError as err:
+            if err.errno != errno.EEXIST or os.path.isfile(directory):
+                raise
+
+
+def run(*args, **kwargs):
+    """Run the command with the given arguments.
+
+    The first argument is the path to the command to run.
+    Subsequent arguments are command-line arguments to be passed.
+
+    This function accepts all optional keyword arguments accepted by
+    `subprocess.Popen`.
+    """
+    args = [i for i in args if i is not None]
+    pipe = subprocess.PIPE
+    process = subprocess.Popen(
+        args, stdout=kwargs.pop('stdout', pipe),
+        stderr=kwargs.pop('stderr', pipe),
+        close_fds=kwargs.pop('close_fds', True), **kwargs)
+    stdout, stderr = process.communicate()
+    if process.returncode:
+        raise subprocess.CalledProcessError(
+            process.returncode, repr(args), output=stdout+stderr)
+    return stdout
+
+
+def script_name():
+    """Return the name of this script."""
+    return os.path.basename(sys.argv[0])
+
+
+def search_file(regexp, filename):
+    """Return the first line in `filename` that matches `regexp`."""
     with open(filename) as f:
         for line in f:
-            if re.match(content, line):
-                return line.strip()
-
-
-def get_value_from_line(line):
-    return line.split('=')[1].strip('"\' ')
-
-
-def script_name():
-    return os.path.basename(sys.argv[0])
+            if re.search(regexp, line):
+                return line
+
+
+def ssh(location, user=None, key=None, caller=subprocess.call):
+    """Return a callable that can be used to run ssh shell commands.
+
+    The ssh `location` and, optionally, `user` must be given.
+    If the user is None then the current user is used for the connection.
+
+    The callable internally uses the given `caller`::
+
+        >>> def caller(cmd):
+        ...     print tuple(cmd)
+        >>> sshcall = ssh('example.com', 'myuser', caller=caller)
+        >>> root_sshcall = ssh('example.com', caller=caller)
+        >>> sshcall('ls -l') # doctest: +ELLIPSIS
+        ('ssh', '-t', ..., '[e-mail]', '--', 'ls -l')
+        >>> root_sshcall('ls -l') # doctest: +ELLIPSIS
+        ('ssh', '-t', ..., 'example.com', '--', 'ls -l')
+
+    The ssh key path can be optionally provided::
+
+        >>> root_sshcall = ssh('example.com', key='/tmp/foo', caller=caller)
+        >>> root_sshcall('ls -l') # doctest: +ELLIPSIS
+        ('ssh', '-t', ..., '-i', '/tmp/foo', 'example.com', '--', 'ls -l')
+
+    If the ssh command exits with an error code,
+    a `subprocess.CalledProcessError` is raised::
+
+        >>> ssh('loc', caller=lambda cmd: 1)('ls -l') # doctest: +ELLIPSIS
+        Traceback (most recent call last):
+        CalledProcessError: ...
+
+    If ignore_errors is set to True when executing the command, no error
+    will be raised, even if the command itself returns an error code.
+
+        >>> sshcall = ssh('loc', caller=lambda cmd: 1)
+        >>> sshcall('ls -l', ignore_errors=True)
+    """
+    sshcmd = [
+        'ssh',
+        '-t',
+        '-t',  # Yes, this second -t is deliberate. See `man ssh`.
+        '-o', 'StrictHostKeyChecking=no',
+        '-o', 'UserKnownHostsFile=/dev/null',
+        ]
+    if key is not None:
+        sshcmd.extend(['-i', key])
+    if user is not None:
+        location = '{}@{}'.format(user, location)
+    sshcmd.extend([location, '--'])
+
+    def _sshcall(cmd, ignore_errors=False):
+        command = sshcmd + [cmd]
+        retcode = caller(command)
+        if retcode and not ignore_errors:
+            raise subprocess.CalledProcessError(retcode, ' '.join(command))
+
+    return _sshcall
+
+
+@contextmanager
+def su(user):
+    """A context manager to temporarily run the script as a different user."""
+    uid, gid = get_user_ids(user)
+    os.setegid(gid)
+    os.seteuid(uid)
+    home = get_user_home(user)
+    with environ(HOME=home):
+        try:
+            yield Env(uid, gid, home)
+        finally:
+            os.setegid(os.getgid())
+            os.seteuid(os.getuid())
+
+
+def user_exists(username):
+    """Return True if given `username` exists, e.g.::
+
+        >>> user_exists('root')
+        True
+        >>> user_exists('_this_user_does_not_exist_')
+        False
+    """
+    try:
+        pwd.getpwnam(username)
+    except KeyError:
+        return False
+    return True
 
 
 def wait_for_page_contents(url, contents, timeout=120, validate=None):
@@ -148,83 +541,6 @@
         time.sleep(0.1)
 
 
-class Serializer:
-    """Handle JSON (de)serialization."""
-
-    def __init__(self, path, default=None, serialize=None, deserialize=None):
-        self.path = path
-        self.default = default or {}
-        self.serialize = serialize or json.dump
-        self.deserialize = deserialize or json.load
-
-    def exists(self):
-        return os.path.exists(self.path)
-
-    def get(self):
-        if self.exists():
-            with open(self.path) as f:
-                return self.deserialize(f)
-        return self.default
-
-    def set(self, data):
-        with open(self.path, 'w') as f:
-            self.serialize(data, f)
-
-
-def get_user_ids(user):
-    """Return the uid and gid of given `user`, e.g.::
-
-        >>> get_user_ids('root')
-        (0, 0)
-    """
-    userdata = pwd.getpwnam(user)
-    return userdata.pw_uid, userdata.pw_gid
-
-
-def get_user_home(user):
-    """Return the home directory of the given `user`.
-
-        >>> get_user_home('root')
-        '/root'
-    """
-    return pwd.getpwnam(user).pw_dir
-
-
-@contextmanager
-def cd(directory):
-    """A context manager to temporary change current working dir, e.g.::
-
-        >>> import os
-        >>> os.chdir('/tmp')
-        >>> with cd('/bin'): print os.getcwd()
-        /bin
-        >>> os.getcwd()
-        '/tmp'
-    """
-    cwd = os.getcwd()
-    os.chdir(directory)
-    yield
-    os.chdir(cwd)
-
-
-@contextmanager
-def su(user):
-    """A context manager to temporary run the script as a different user."""
-    uid, gid = get_user_ids(user)
-    os.setegid(gid)
-    os.seteuid(uid)
-    current_home = os.getenv('HOME')
-    home = get_user_home(user)
-    os.environ['HOME'] = home
-    try:
-        yield Env(uid, gid, home)
-    finally:
-        os.setegid(os.getgid())
-        os.seteuid(os.getuid())
-        if current_home is not None:
-            os.environ['HOME'] = current_home
-
-
 class DictDiffer:
     """
     Calculate the difference between two dictionaries as:

=== modified file 'tests.py'
--- tests.py	2012-03-02 13:58:36 +0000
+++ tests.py	2012-03-07 10:19:06 +0000
@@ -6,44 +6,42 @@
 __metaclass__ = type
 
 
+import getpass
 import os
 from subprocess import CalledProcessError
+import tempfile
 import unittest
 
 from shelltoolbox import (
     cd,
     command,
     DictDiffer,
+    environ,
+    file_append,
+    file_prepend,
+    generate_ssh_keys,
+    get_su_command,
+    get_user_home,
+    get_user_ids,
+    join_command,
+    mkdirs,
     run,
+    search_file,
+    Serializer,
+    ssh,
     su,
+    user_exists,
     )
 
 
-class TestRun(unittest.TestCase):
-
-    def testSimpleCommand(self):
-        # Running a simple command (ls) works and running the command
-        # produces a string.
-        self.assertIsInstance(run('/bin/ls'), str)
-
-    def testStdoutReturned(self):
-        # Running a simple command (ls) works and running the command
-        # produces a string.
-        self.assertIn('Usage:', run('/bin/ls', '--help'))
-
-    def testCalledProcessErrorRaised(self):
-        # If an error occurs a CalledProcessError is raised with the return
-        # code, command executed, and the output of the command.
-        with self.assertRaises(CalledProcessError) as info:
-            run('ls', '--not a valid switch')
-        exception = info.exception
-        self.assertEqual(2, exception.returncode)
-        self.assertEqual("['ls', '--not a valid switch']", exception.cmd)
-        self.assertIn('unrecognized option', exception.output)
-
-    def testNoneArguments(self):
-        # Ensure None is ignored when passed as positional argument.
-        self.assertIn('Usage:', run('/bin/ls', None, '--help', None))
+class TestCdContextManager(unittest.TestCase):
+
+    def test_cd(self):
+        curdir = os.getcwd()
+        self.assertNotEqual('/var', curdir)
+        with cd('/var'):
+            self.assertEqual('/var', os.getcwd())
+        self.assertEqual(curdir, os.getcwd())
 
 
 class TestCommand(unittest.TestCase):
@@ -112,6 +110,349 @@
         self.assertEquals(expected, diff.added_or_changed)
 
 
+class TestEnviron(unittest.TestCase):
+
+    def test_existing(self):
+        # If an existing environment variable is changed, it is
+        # restored during context cleanup.
+        os.environ['MY_VARIABLE'] = 'foo'
+        with environ(MY_VARIABLE='bar'):
+            self.assertEqual('bar', os.getenv('MY_VARIABLE'))
+        self.assertEqual('foo', os.getenv('MY_VARIABLE'))
+        del os.environ['MY_VARIABLE']
+
+    def test_new(self):
+        # If a new environment variable is added, it is removed during
+        # context cleanup.
+        with environ(MY_VAR1='foo', MY_VAR2='bar'):
+            self.assertEqual('foo', os.getenv('MY_VAR1'))
+            self.assertEqual('bar', os.getenv('MY_VAR2'))
+        self.assertIsNone(os.getenv('MY_VAR1'))
+        self.assertIsNone(os.getenv('MY_VAR2'))
+
+
+class BaseCreateFile(object):
+
+    def create_file(self, content):
+        f = tempfile.NamedTemporaryFile('w', delete=False)
+        f.write(content)
+        f.close()
+        return f
+
+
+class BaseTestFile(BaseCreateFile):
+
+    base_content = 'line1\n'
+    new_content = 'new line\n'
+
+    def check_file_content(self, content, filename):
+        self.assertEqual(content, open(filename).read())
+
+
+class TestFileAppend(unittest.TestCase, BaseTestFile):
+
+    def test_append(self):
+        # Ensure the new line is correctly added at the end of the file.
+        f = self.create_file(self.base_content)
+        file_append(f.name, self.new_content)
+        self.check_file_content(self.base_content + self.new_content, f.name)
+
+    def test_existing_content(self):
+        # Ensure nothing happens if the file already contains the given line.
+        content = self.base_content + self.new_content
+        f = self.create_file(content)
+        file_append(f.name, self.new_content)
+        self.check_file_content(content, f.name)
+
+    def test_new_line_in_file_contents(self):
+        # A new line is automatically added before the given content if it
+        # is not present at the end of current file.
+        f = self.create_file(self.base_content.strip())
+        file_append(f.name, self.new_content)
+        self.check_file_content(self.base_content + self.new_content, f.name)
+
+    def test_new_line_in_given_line(self):
+        # A new line is automatically added to the given line if not present.
+        f = self.create_file(self.base_content)
+        file_append(f.name, self.new_content.strip())
+        self.check_file_content(self.base_content + self.new_content, f.name)
+
+    def test_non_existent_file(self):
+        # Ensure the file is created if it does not exist.
+        filename = tempfile.mktemp()
+        file_append(filename, self.base_content)
+        self.check_file_content(self.base_content, filename)
+
+    def test_fragment(self):
+        # Ensure a line fragment is not matched.
+        f = self.create_file(self.base_content)
+        fragment = self.base_content[2:]
+        file_append(f.name, fragment)
+        self.check_file_content(self.base_content + fragment, f.name)
+
+
+class TestFilePrepend(unittest.TestCase, BaseTestFile):
+
+    def test_prpend(self):
+        # Ensure the new content is correctly prepended at the beginning of
+        # the file.
+        f = self.create_file(self.base_content)
+        file_prepend(f.name, self.new_content)
+        self.check_file_content(self.new_content + self.base_content, f.name)
+
+    def test_existing_content(self):
+        # Ensure nothing happens if the file already starts with the given
+        # content.
+        content = self.base_content + self.new_content
+        f = self.create_file(content)
+        file_prepend(f.name, self.base_content)
+        self.check_file_content(content, f.name)
+
+    def test_move_content(self):
+        # If the file contains the given content, but not at the beginning,
+        # the content is moved on top.
+        f = self.create_file(self.base_content + self.new_content)
+        file_prepend(f.name, self.new_content)
+        self.check_file_content(self.new_content + self.base_content, f.name)
+
+    def test_new_line_in_given_line(self):
+        # A new line is automatically added to the given line if not present.
+        f = self.create_file(self.base_content)
+        file_prepend(f.name, self.new_content.strip())
+        self.check_file_content(self.new_content + self.base_content, f.name)
+
+
+class TestGenerateSSHKeys(unittest.TestCase):
+
+    def test_generation(self):
+        # Ensure ssh keys are correctly generated.
+        filename = tempfile.mktemp()
+        generate_ssh_keys(filename)
+        first_line = open(filename).readlines()[0].strip()
+        self.assertEqual('-----BEGIN RSA PRIVATE KEY-----', first_line)
+        pub_content = open(filename + '.pub').read()
+        self.assertTrue(pub_content.startswith('ssh-rsa'))
+
+
+class TestGetSuCommand(unittest.TestCase):
+
+    def test_current_user(self):
+        # If the su is requested as current user, the arguments are
+        # returned as given.
+        cmd = ('ls', '-l')
+        command = get_su_command(getpass.getuser(), cmd)
+        self.assertSequenceEqual(cmd, command)
+
+    def test_another_user(self):
+        # Ensure "su" is prepended and arguments are correctly quoted.
+        command = get_su_command('nobody', ('ls', '-l', 'my file'))
+        self.assertSequenceEqual(
+            ('su', 'nobody', '-c', "ls -l 'my file'"), command)
+
+
+class TestGetUserHome(unittest.TestCase):
+
+    def test_existent(self):
+        # Ensure the real home directory is returned for existing users.
+        self.assertEqual('/root', get_user_home('root'))
+
+    def test_non_existent(self):
+        # If the user does not exist, return a default /home/[username] home.
+        user = '_this_user_does_not_exist_'
+        self.assertEqual('/home/' + user, get_user_home(user))
+
+
+class TestGetUserIds(unittest.TestCase):
+
+    def test_get_user_ids(self):
+        # Ensure the correct uid and gid are returned.
+        uid, gid = get_user_ids('root')
+        self.assertEqual(0, uid)
+        self.assertEqual(0, gid)
+
+
+class TestJoinCommand(unittest.TestCase):
+
+    def test_normal(self):
+        # Ensure a normal command is correctly parsed.
+        command = 'ls -l'
+        self.assertEqual(command, join_command(command.split()))
+
+    def test_containing_spaces(self):
+        # Ensure args containing spaces are correctly quoted.
+        args = ('command', 'arg containig spaces')
+        self.assertEqual("command 'arg containig spaces'", join_command(args))
+
+    def test_empty(self):
+        # Ensure empty args are correctly quoted.
+        args = ('command', '')
+        self.assertEqual("command ''", join_command(args))
+
+
+class TestMkdirs(unittest.TestCase):
+
+    def test_intermediate_dirs(self):
+        # Ensure the leaf directory and all intermediate ones are created.
+        base_dir = tempfile.mktemp(suffix='/')
+        dir1 = tempfile.mktemp(prefix=base_dir)
+        dir2 = tempfile.mktemp(prefix=base_dir)
+        mkdirs(dir1, dir2)
+        self.assertTrue(os.path.isdir(dir1))
+        self.assertTrue(os.path.isdir(dir2))
+
+    def test_existing_dir(self):
+        # If the leaf directory already exists the function returns
+        # without errors.
+        mkdirs('/tmp')
+
+    def test_existing_file(self):
+        # An `OSError` is raised if the leaf path exists and it is a file.
+        f = tempfile.NamedTemporaryFile('w', delete=False)
+        f.close()
+        with self.assertRaises(OSError):
+            mkdirs(f.name)
+
+
+class TestRun(unittest.TestCase):
+
+    def testSimpleCommand(self):
+        # Running a simple command (ls) works and running the command
+        # produces a string.
+        self.assertIsInstance(run('/bin/ls'), str)
+
+    def testStdoutReturned(self):
+        # Running a simple command (ls) works and running the command
+        # produces a string.
+        self.assertIn('Usage:', run('/bin/ls', '--help'))
+
+    def testCalledProcessErrorRaised(self):
+        # If an error occurs a CalledProcessError is raised with the return
+        # code, command executed, and the output of the command.
+        with self.assertRaises(CalledProcessError) as info:
+            run('ls', '--not a valid switch')
+        exception = info.exception
+        self.assertEqual(2, exception.returncode)
+        self.assertEqual("['ls', '--not a valid switch']", exception.cmd)
+        self.assertIn('unrecognized option', exception.output)
+
+    def testNoneArguments(self):
+        # Ensure None is ignored when passed as positional argument.
+        self.assertIn('Usage:', run('/bin/ls', None, '--help', None))
+
+
+class TestSearchFile(unittest.TestCase, BaseCreateFile):
+
+    content1 = 'content1\n'
+    content2 = 'content2\n'
+
+    def setUp(self):
+        self.filename = self.create_file(self.content1 + self.content2).name
+
+    def tearDown(self):
+        os.remove(self.filename)
+
+    def test_grep(self):
+        # Ensure plain text is correctly matched.
+        self.assertEqual(self.content2, search_file('ent2', self.filename))
+        self.assertEqual(self.content1, search_file('content', self.filename))
+
+    def test_no_match(self):
+        # Ensure the function does not return false positives.
+        self.assertIsNone(search_file('no_match', self.filename))
+
+    def test_regexp(self):
+        # Ensure the function works with regular expressions.
+        self.assertEqual(self.content2, search_file('\w2', self.filename))
+
+
+class TestSerializer(unittest.TestCase):
+
+    def setUp(self):
+        self.path = tempfile.mktemp()
+        self.data = {'key': 'value'}
+
+    def tearDown(self):
+        if os.path.exists(self.path):
+            os.remove(self.path)
+
+    def test_serializer(self):
+        # Ensure data is correctly serializied and deserialized.
+        s = Serializer(self.path)
+        s.set(self.data)
+        self.assertEqual(self.data, s.get())
+
+    def test_existence(self):
+        # Ensure the file is created only when needed.
+        s = Serializer(self.path)
+        self.assertFalse(s.exists())
+        s.set(self.data)
+        self.assertTrue(s.exists())
+
+    def test_default_value(self):
+        # If the file does not exist, the serializer returns a default value.
+        s = Serializer(self.path)
+        self.assertEqual({}, s.get())
+        s = Serializer(self.path, default=47)
+        self.assertEqual(47, s.get())
+
+    def test_another_serializer(self):
+        # It is possible to use a custom serializer (e.g. pickle).
+        import pickle
+        s = Serializer(
+            self.path, serialize=pickle.dump, deserialize=pickle.load)
+        s.set(self.data)
+        self.assertEqual(self.data, s.get())
+
+
+class TestSSH(unittest.TestCase):
+
+    def setUp(self):
+        self.last_command = None
+
+    def remove_command_options(self, cmd):
+        cmd = list(cmd)
+        del cmd[1:7]
+        return cmd
+
+    def caller(self, cmd):
+        self.last_command = self.remove_command_options(cmd)
+
+    def check_last_command(self, expected):
+        self.assertSequenceEqual(expected, self.last_command)
+
+    def test_current_user(self):
+        # Ensure ssh command is correctly generated for current user.
+        sshcall = ssh('example.com', caller=self.caller)
+        sshcall('ls -l')
+        self.check_last_command(['ssh', 'example.com', '--', 'ls -l'])
+
+    def test_another_user(self):
+        # Ensure ssh command is correctly generated for a different user.
+        sshcall = ssh('example.com', 'myuser', caller=self.caller)
+        sshcall('ls -l')
+        self.check_last_command(['ssh', '[e-mail]', '--', 'ls -l'])
+
+    def test_ssh_key(self):
+        # The ssh key path can be optionally provided.
+        sshcall = ssh('example.com', key='/tmp/foo', caller=self.caller)
+        sshcall('ls -l')
+        self.check_last_command([
+            'ssh', '-i', '/tmp/foo', 'example.com', '--', 'ls -l'])
+
+    def test_error(self):
+        # If the ssh command exits with an error code, a
+        # `subprocess.CalledProcessError` is raised.
+        sshcall = ssh('example.com', caller=lambda cmd: 1)
+        with self.assertRaises(CalledProcessError):
+            sshcall('ls -l')
+
+    def test_ignore_errors(self):
+        # If ignore_errors is set to True when executing the command, no error
+        # will be raised, even if the command itself returns an error code.
+        sshcall = ssh('example.com', caller=lambda cmd: 1)
+        sshcall('ls -l', ignore_errors=True)
+
+
 current_euid = os.geteuid()
 current_egid = os.getegid()
 current_home = os.environ['HOME']
@@ -180,13 +521,11 @@
             self.assertEqual(current_home, os.environ['HOME'])
 
 
-class TestCdContextManager(unittest.TestCase):
-    def test_cd(self):
-        curdir = os.getcwd()
-        self.assertNotEqual('/var', curdir)
-        with cd('/var'):
-            self.assertEqual('/var', os.getcwd())
-        self.assertEqual(curdir, os.getcwd())
+class TestUserExists(unittest.TestCase):
+
+    def test_user_exists(self):
+        self.assertTrue(user_exists('root'))
+        self.assertFalse(user_exists('_this_user_does_not_exist_'))
 
 
 if __name__ == '__main__':