tasks.py

# coding: utf-8
from __future__ import unicode_literals
from __future__ import absolute_import
from invoke import run
from invoke import task


import io
import json
import os
import re


PROJECT_ROOT_DIR = os.path.dirname(os.path.abspath(__file__))

ANSIBLE_CFG = '''
[ssh_connection]
control_path = %(directory)s/%%h-%%p-%%r
scp_if_ssh = True

[defaults]
inventory = ./inventory
gathering = smart
fact_caching_timeout = 300
host_key_checking = False
retry_files_enabled = False
fact_caching = jsonfile
ansible_managed = Ansible managed: Do NOT edit this file manually!
nocows = 1
pipelining = True
fact_caching_connection = tmp/ansible-facts
hash_behaviour = merge
'''


def rm_f(path):
    if os.path.isfile(path):
        os.unlink(path)
        return True


def touch(filename):
    """Simulate touch

    """
    with io.open(filename, mode='a', encoding='utf-8'):
        pass


def vagrant_status():
    """Get status of vagrant box

    Returns:
        `generator` - (box `str`, status `str`, provider `str`

    """
    result = run('vagrant status', hide=True)
    out = result.stdout
    assert result.ok, 'Unable to get list of vagrant environments, please invoke `vagrant status`'

    for x in out.splitlines():
        regex = re.compile(r'\([a-zA-Z0-9]+\)')
        x = x.encode().decode()
        if regex.search(x):
            x = x.replace(')', '')
            env, status, provider = re.split(r' \(|  +', x)
            yield (env, status, provider)


def vagrant_halt(box='default', **kwargs):
    """Poweroff a vagrant box

    """
    return run(
        'vagrant halt --force {box}'.format(box=box),
        **kwargs
    )


def vagrant_up(box='default', **kwargs):
    """Starts a vagrant box

    """
    # remove state file
    rm_f_built_state(box=box)

    result = None

    if is_active(box=box) is False:
        result = run(
            'vagrant up {box}'.format(box=box),
            **kwargs
        )

    # create the built state file
    touch_built_state(box=box)

    return result


def vagrant_destroy(box='default', **kwargs):
    """Destroys a vagrant box

    """
    return run(
        'vagrant destroy --force {box}'.format(box=box),
        **kwargs
    )


def vagrant_provision(box='default', **kwargs):
    """Run the vagrant provisioner

    """
    return run(
        'vagrant provision {box}'.format(box=box),
        **kwargs
    )


def vagrant_snapshot_list(name, box='default', **kwargs):
    """Return a list of vagrant box snapshots

    """
    result = run(
        'vagrant snapshot list {box}'.format(box=box),
        hide=True,
        **kwargs
    )
    out = result.stdout
    result = [line.replace('==>', '').replace('default:', '').strip() for line in out.splitlines()]
    return result


def vagrant_snapshot(name, box='default', **kwargs):
    """Creates a new snapshot, if name does not already exist

    Checks to see if the snapshot name for the box already exists?
      - if no snapshot name exists
         - create a new snapshot called {name} for {box}
      - if snapshot name exists
         - do nothing

    """
    if has_snapshot(box=box, name=name, **kwargs) is False:
        return run(
            'vagrant snapshot save {box} {snapshot}'.format(box=box, snapshot=name),
            hide=True,
            **kwargs
        )


def vagrant_ssh_config(box='default', filename=None, *args, **kwargs):
    """Returns the ssh-config for box

    """
    args = ' '.join(args)

    if is_active(box=box) is False:
        return None

    # otherwise the box is active so we can get at the ssh config
    cmd = 'vagrant ssh-config {box} {args}'.format(box=box, args=args)
    result = run(cmd, **kwargs)

    assert result.return_code == 0, 'Expected command: `{}` to return exit code: 0 got: {}'.format(cmd, result.return_code)

    ssh_config = result.stdout.decode('utf-8').strip() or result.stderr.decode('utf-8').strip()

    if filename:
        with io.open(filename, encoding='utf-8', mode='w') as fd:
            fd.write(ssh_config)
            fd.write('\n')

    return ssh_config


def run_converger(box='default', **kwargs):
    """Converge the vagrant environment

    """
    # remove converged state file
    rm_f_converged_state(box=box)

    result = run(
        'ansible-playbook converge.yml -vv -l {box}'.format(box=box),
        **kwargs
    )

    # create the converged state file
    touch_converged_state(box=box)

    return result


def run_verifier(*args, **kwargs):
    """Runs the verifier

    """
    default_args = (
        'tox',
    )

    cmd = default_args + args
    return run(' '.join(cmd), **kwargs)


def has_snapshot(name, box='default', **kwargs):
    """Returns `True` if snapshot name exists, otherwise returns `False`

    """
    # we remove pty becuase we don't want the color codes mixed into the output
    kwargs['pty'] = False
    out = vagrant_snapshot_list(name=name, box=box, **kwargs)

    return name in out


def is_built(box):
    """Returns `True` if box is already built, otherwise returns `False`

    """
    return os.path.isfile('.{box}.built'.format(box=box))


def touch_built_state(box):
    """Create the box built state file

    """
    filename = os.path.join(PROJECT_ROOT_DIR, '.{box}.built'.format(box=box))
    return touch(filename)


def rm_f_built_state(box):
    """Remove the box built state file

    """
    filename = os.path.join(PROJECT_ROOT_DIR, '.{box}.built'.format(box=box))
    return rm_f(filename)


def is_converged(box):
    """Returns `True` if box is already converged, otherwise returns `False`

    """
    return os.path.isfile('.{box}.converged'.format(box=box))


def touch_converged_state(box):
    """Create the box converged state file

    """
    filename = os.path.join(PROJECT_ROOT_DIR, '.{box}.converged'.format(box=box))
    return touch(filename)


def rm_f_converged_state(box):
    """Remove the box converged state file

    """
    filename = os.path.join(PROJECT_ROOT_DIR, '.{box}.converged'.format(box=box))
    return rm_f(filename)


def is_active(box, provider='virtualbox', **kwargs):
    """Returns `True` if box is active, otherwise returns `False`

    """
    off_states = (
        'aborted',
        'frozen',
        'not created',
        'poweroff',
        'shutoff',
        'paused',
        'saved',
    )

    for env, status, provider in vagrant_status():
        if provider == provider and env == box and status in off_states:
            return False

        return True

    return False


def get_vagrant_suts():
    """Return a `generator` of defined vagrant test environments

    """
    for sut, status, provider in vagrant_status():
        yield dict(sut=sut,
                   status=status,
                   provider=provider,
                   built=is_built(box=sut),
                   converged=is_converged(box=sut),
                   )


@task
def clean(c, echo=False):
    """Cleans all compiled artifacts recursively

    """
    file_patterns = [
        '*.pyc',
        '*.retry',
        '.*.built',
        '.*.converged',
        '.coverage',
    ]

    dir_patterns = [
        'build',
        'dist',
        '*.egg-info',
        '.tox',
        '__pycache__',
        'tmp',
    ]

    for pattern in file_patterns:
        c.run('find . -name "{}" -type f | xargs rm -f'.format(pattern), echo=echo)

    for pattern in dir_patterns:
        c.run('find . -name "{}" -type d | xargs rm -rf'.format(pattern), echo=echo)


@task
def coverage(c, extra=None, pty=True):
    """Run code coverage

    """
    if os.path.isfile('.coverage'):
        os.unlink('.coverage')

    default_args = [
        'coverage',
        'run',
        '-m',
        'pytest',
        '-v',
        '-rs',
        '-s',
    ]

    cmd = ' '.join(default_args)
    if extra:
        cmd += ' -- {}'.format(extra)
    c.run(cmd, pty=pty)
    c.run('coverage report -m', pty=pty)


@task
def prep(c, pty=True, box='default', skip_if_built=False):
    """Create the VM, Power it on, run the prep provisioner

    """
    snapshot = 'prep'
    if skip_if_built and is_built(box=box):
        return None

    vagrant_up(box=box, pty=pty)
    vagrant_provision(box=box, pty=pty)
    vagrant_snapshot(name=snapshot, box=box, pty=pty)


@task
def list(c, pty=True):
    """Lists test environments

    """
    envs = [x for x in get_vagrant_suts()]
    sut_json = json.dumps(envs, sort_keys=True, indent=4, ensure_ascii=False)
    print(sut_json)


@task
def converge(c, box='default', pty=True, skip_if_built=True, skip_if_converged=False):
    """Runs the custom code that puts the VM into the desired state

    """
    snapshot = 'converged'

    if skip_if_converged and is_converged(box=box):
        return None

    prep(c, box=box, pty=pty, skip_if_built=skip_if_built)
    run_converger(box=box, pty=pty)
    vagrant_snapshot(name=snapshot, box=box, pty=pty)


@task
def test(c, box='default', pty=True, skip_if_built=True, skip_if_converged=True, verify=True, destroy_box=False, tox_args=[]):
    """Starts the VM if not running, converge, runs tests, shuts the VM down

    """
    prep(c, box=box, pty=pty, skip_if_built=skip_if_built)
    converge(c, box=box, pty=pty, skip_if_built=skip_if_built, skip_if_converged=skip_if_converged)

    if verify:
        run_verifier(*tox_args, pty=pty)

    if destroy_box:
        destroy(c, box=box, pty=pty)


@task
def destroy(c, box='default', pty=True):
    """Destroys a vagrant box

    """
    rm_f_built_state(box=box)
    rm_f_converged_state(box=box)
    vagrant_destroy(box=box, pty=pty)


@task
def ssh_config(c, box='default', filename=None, pty=True):
    """Returns the ssh-config for the box if it is running

    """
    vagrant_ssh_config(box=box, filename=filename, pty=pty)


@task
def ansible_inventory(c, box='default', pty=True):
    """Creates the inventory file used for ansible

    """
    filename = os.path.join(PROJECT_ROOT_DIR, 'inventory')
    ssh_config = vagrant_ssh_config(box=box, filename='ssh-config.{}'.format(box), hide=True)
    ssh_config = dict([re.sub('\s\s+', '', x).split(' ', 1) for x in ssh_config.splitlines()])

    inventory = (
        '{box} '
        'ansible_host={ansible_host} '
        'ansible_port={ansible_port} '
        'ansible_user={ansible_user} '
        'ansible_ssh_private_key_file={ansible_ssh_private_key_file}'
    ).format(
        box=box,
        ansible_host=ssh_config['HostName'],
        ansible_port=ssh_config['Port'],
        ansible_user=ssh_config['User'],
        ansible_ssh_private_key_file=ssh_config['IdentityFile']
    )

    with io.open(filename, encoding='utf-8', mode='w') as fd:
        fd.write(inventory)
        fd.write("\n")