<commit_before># Analyzes an object and outputs numeric properties

import cv2
import numpy as np
import os
from plantcv.plantcv import fatal_error
from plantcv.plantcv import print_image
from plantcv.plantcv import plot_image
from plantcv.plantcv import rgb2gray_hsv
from plantcv.plantcv import find_objects
from plantcv.plantcv.threshold import binary as binary_threshold
from plantcv.plantcv import roi_objects
from plantcv.plantcv import object_composition
from plantcv.plantcv import apply_mask
from plantcv.plantcv import params

def report_size_marker_area(img, roi_contour, roi_hierarchy, marker='define', objcolor='dark', thresh_channel=None,
                            thresh=None, filename=False):
    """Detects a size marker in a specified region and reports its size and eccentricity

    Inputs:
    img             = An RGB or grayscale image to plot the marker object on
    roi_contour     = A region of interest contour (e.g. output from pcv.roi.rectangle or other methods)
    roi_hierarchy   = A region of interest contour hierarchy (e.g. output from pcv.roi.rectangle or other methods)
    marker          = 'define' or 'detect'. iIf define it means you set an area, if detect it means you want to
                      detect within an area
    objcolor        = Object color is 'dark' or 'light' (is the marker darker or lighter than the background)
    thresh_channel  = 'h', 's', or 'v' for hue, saturation or value
    thresh          = Binary threshold value (integer)
    filename        = False or the name of an output image file

    Returns:
    marker_header   = Marker data table headers
    marker_data     = Marker data table values
    analysis_images = List of output images

    :param img: numpy.ndarray
    :param roi_contour: list
    :param roi_hierarchy: numpy.ndarray
    :param marker: str
    :param objcolor: str
    :param thresh_channel: str
    :param thresh: int
    :param filename: str
    :return: marker_header: list
    :return: marker_data: list
    :return: analysis_images: list
    """

    params.device += 1
    # Make a copy of the reference image
    ref_img = np.copy(img)
    # If the reference image is grayscale convert it to color
    if len(np.shape(ref_img)) == 2:
        ref_img = cv2.cvtColor(ref_img, cv2.COLOR_GRAY2BGR)

    # Marker components
    # If the marker type is "defined" then the marker_mask and marker_contours are equal to the input ROI
    # Initialize a binary image
    roi_mask = np.zeros(np.shape(img)[:2], dtype=np.uint8)
    # Draw the filled ROI on the mask
    cv2.drawContours(roi_mask, roi_contour, -1, (255), -1)
    marker_mask = []
    marker_contour = []

    # If the marker type is "detect" then we will use the ROI to isolate marker contours from the input image
    if marker == 'detect':
        # We need to convert the input image into an one of the HSV channels and then threshold it
        if thresh_channel is not None and thresh is not None:
            # Mask the input image
            masked = apply_mask(img=ref_img, mask=roi_mask, mask_color="black")
            # Convert the masked image to hue, saturation, or value
            marker_hsv = rgb2gray_hsv(img=masked, channel=thresh_channel)
            # Threshold the HSV image
            marker_bin = binary_threshold(img=marker_hsv, threshold=thresh, maxValue=255, object_type=objcolor)
            # Identify contours in the masked image
            contours, hierarchy = find_objects(img=ref_img, mask=marker_bin)
            # Filter marker contours using the input ROI
            kept_contours, kept_hierarchy, kept_mask, obj_area = roi_objects(img=ref_img, object_contour=contours,
                                                                             obj_hierarchy=hierarchy,
                                                                             roi_contour=roi_contour,
                                                                             roi_hierarchy=roi_hierarchy,
                                                                             roi_type="partial")
            # If there are more than one contour detected, combine them into one
            # These become the marker contour and mask
            marker_contour, marker_mask = object_composition(img=ref_img, contours=kept_contours,
                                                             hierarchy=kept_hierarchy)
        else:
            fatal_error('thresh_channel and thresh must be defined in detect mode')
    elif marker == "define":
        # Identify contours in the masked image
        contours, hierarchy = find_objects(img=ref_img, mask=roi_mask)
        # If there are more than one contour detected, combine them into one
        # These become the marker contour and mask
        marker_contour, marker_mask = object_composition(img=ref_img, contours=contours, hierarchy=hierarchy)
    else:
        fatal_error("marker must be either 'define' or 'detect' but {0} was provided.".format(marker))

    # Calculate the moments of the defined marker region
    m = cv2.moments(marker_mask, binaryImage=True)
    # Calculate the marker area
    marker_area = m['m00']

    # Fit a bounding ellipse to the marker
    center, axes, angle = cv2.fitEllipse(marker_contour)
    major_axis = np.argmax(axes)
    minor_axis = 1 - major_axis
    major_axis_length = axes[major_axis]
    minor_axis_length = axes[minor_axis]
    # Calculate the bounding ellipse eccentricity
    eccentricity = np.sqrt(1 - (axes[minor_axis] / axes[major_axis]) ** 2)

    # Make a list to store output images
    analysis_images = []
    cv2.drawContours(ref_img, marker_contour, -1, (255, 0, 0), 5)
    if filename:
        out_file = str(filename[0:-4]) + '_sizemarker.jpg'
        print_image(ref_img, out_file)
        analysis_images.append(['IMAGE', 'marker', out_file])
    if params.debug is 'print':
        print_image(ref_img, os.path.join(params.debug_outdir, str(params.device) + '_marker_shape.png'))
    elif params.debug is 'plot':
        plot_image(ref_img)

    marker_header = (
        'HEADER_MARKER',
        'marker_area',
        'marker_major_axis_length',
        'marker_minor_axis_length',
        'marker_eccentricity'
    )

    marker_data = (
        'MARKER_DATA',
        marker_area,
        major_axis_length,
        minor_axis_length,
        eccentricity
    )

    return marker_header, marker_data, analysis_images
<commit_msg>Fixed incorrect input name<commit_after># Analyzes an object and outputs numeric properties

import cv2
import numpy as np
import os
from plantcv.plantcv import fatal_error
from plantcv.plantcv import print_image
from plantcv.plantcv import plot_image
from plantcv.plantcv import rgb2gray_hsv
from plantcv.plantcv import find_objects
from plantcv.plantcv.threshold import binary as binary_threshold
from plantcv.plantcv import roi_objects
from plantcv.plantcv import object_composition
from plantcv.plantcv import apply_mask
from plantcv.plantcv import params

def report_size_marker_area(img, roi_contour, roi_hierarchy, marker='define', objcolor='dark', thresh_channel=None,
                            thresh=None, filename=False):
    """Detects a size marker in a specified region and reports its size and eccentricity

    Inputs:
    img             = An RGB or grayscale image to plot the marker object on
    roi_contour     = A region of interest contour (e.g. output from pcv.roi.rectangle or other methods)
    roi_hierarchy   = A region of interest contour hierarchy (e.g. output from pcv.roi.rectangle or other methods)
    marker          = 'define' or 'detect'. iIf define it means you set an area, if detect it means you want to
                      detect within an area
    objcolor        = Object color is 'dark' or 'light' (is the marker darker or lighter than the background)
    thresh_channel  = 'h', 's', or 'v' for hue, saturation or value
    thresh          = Binary threshold value (integer)
    filename        = False or the name of an output image file

    Returns:
    marker_header   = Marker data table headers
    marker_data     = Marker data table values
    analysis_images = List of output images

    :param img: numpy.ndarray
    :param roi_contour: list
    :param roi_hierarchy: numpy.ndarray
    :param marker: str
    :param objcolor: str
    :param thresh_channel: str
    :param thresh: int
    :param filename: str
    :return: marker_header: list
    :return: marker_data: list
    :return: analysis_images: list
    """

    params.device += 1
    # Make a copy of the reference image
    ref_img = np.copy(img)
    # If the reference image is grayscale convert it to color
    if len(np.shape(ref_img)) == 2:
        ref_img = cv2.cvtColor(ref_img, cv2.COLOR_GRAY2BGR)

    # Marker components
    # If the marker type is "defined" then the marker_mask and marker_contours are equal to the input ROI
    # Initialize a binary image
    roi_mask = np.zeros(np.shape(img)[:2], dtype=np.uint8)
    # Draw the filled ROI on the mask
    cv2.drawContours(roi_mask, roi_contour, -1, (255), -1)
    marker_mask = []
    marker_contour = []

    # If the marker type is "detect" then we will use the ROI to isolate marker contours from the input image
    if marker == 'detect':
        # We need to convert the input image into an one of the HSV channels and then threshold it
        if thresh_channel is not None and thresh is not None:
            # Mask the input image
            masked = apply_mask(img=ref_img, mask=roi_mask, mask_color="black")
            # Convert the masked image to hue, saturation, or value
            marker_hsv = rgb2gray_hsv(img=masked, channel=thresh_channel)
            # Threshold the HSV image
            marker_bin = binary_threshold(gray_img=marker_hsv, threshold=thresh, max_value=255, object_type=objcolor)
            # Identify contours in the masked image
            contours, hierarchy = find_objects(img=ref_img, mask=marker_bin)
            # Filter marker contours using the input ROI
            kept_contours, kept_hierarchy, kept_mask, obj_area = roi_objects(img=ref_img, object_contour=contours,
                                                                             obj_hierarchy=hierarchy,
                                                                             roi_contour=roi_contour,
                                                                             roi_hierarchy=roi_hierarchy,
                                                                             roi_type="partial")
            # If there are more than one contour detected, combine them into one
            # These become the marker contour and mask
            marker_contour, marker_mask = object_composition(img=ref_img, contours=kept_contours,
                                                             hierarchy=kept_hierarchy)
        else:
            fatal_error('thresh_channel and thresh must be defined in detect mode')
    elif marker == "define":
        # Identify contours in the masked image
        contours, hierarchy = find_objects(img=ref_img, mask=roi_mask)
        # If there are more than one contour detected, combine them into one
        # These become the marker contour and mask
        marker_contour, marker_mask = object_composition(img=ref_img, contours=contours, hierarchy=hierarchy)
    else:
        fatal_error("marker must be either 'define' or 'detect' but {0} was provided.".format(marker))

    # Calculate the moments of the defined marker region
    m = cv2.moments(marker_mask, binaryImage=True)
    # Calculate the marker area
    marker_area = m['m00']

    # Fit a bounding ellipse to the marker
    center, axes, angle = cv2.fitEllipse(marker_contour)
    major_axis = np.argmax(axes)
    minor_axis = 1 - major_axis
    major_axis_length = axes[major_axis]
    minor_axis_length = axes[minor_axis]
    # Calculate the bounding ellipse eccentricity
    eccentricity = np.sqrt(1 - (axes[minor_axis] / axes[major_axis]) ** 2)

    # Make a list to store output images
    analysis_images = []
    cv2.drawContours(ref_img, marker_contour, -1, (255, 0, 0), 5)
    if filename:
        out_file = str(filename[0:-4]) + '_sizemarker.jpg'
        print_image(ref_img, out_file)
        analysis_images.append(['IMAGE', 'marker', out_file])
    if params.debug is 'print':
        print_image(ref_img, os.path.join(params.debug_outdir, str(params.device) + '_marker_shape.png'))
    elif params.debug is 'plot':
        plot_image(ref_img)

    marker_header = (
        'HEADER_MARKER',
        'marker_area',
        'marker_major_axis_length',
        'marker_minor_axis_length',
        'marker_eccentricity'
    )

    marker_data = (
        'MARKER_DATA',
        marker_area,
        major_axis_length,
        minor_axis_length,
        eccentricity
    )

    return marker_header, marker_data, analysis_images
<|endoftext|>


<commit_before>import os
import re
import tempfile

import yaml

from fabric.api import cd, env, execute, get, hide, lcd, local, put, require, run, settings, sudo, task
from fabric.colors import red
from fabric.contrib import files, project
from fabric.contrib.console import confirm
from fabric.utils import abort

DEFAULT_SALT_LOGLEVEL = 'info'
SALT_VERSION = '2015.5.1'
PROJECT_NAME = "service_info"
PROJECT_ROOT = os.path.dirname(__file__)
CONF_ROOT = os.path.join(PROJECT_ROOT, 'conf')

VALID_ROLES = (
    'salt-master',
    'web',
    'worker',
    'balancer',
    'queue',
    'cache',
    'beat',
)

def _common_env():
    env.project = PROJECT_NAME
    # refresh_environment() needs this even though the task is run against staging
    env.production_master = 'ec2-54-93-51-232.eu-central-1.compute.amazonaws.com'

@task
def staging():
    _common_env()
    env.environment = 'staging'
    # 54.93.66.254 Staging server on AWS Frankfurt
    # Internal hostname for our own access to the server
    # To change the domain, see conf/pillar/<envname>/env.sls
    # env.master = 'serviceinfo-staging.rescue.org'
    env.master = 'ec2-54-93-66-254.eu-central-1.compute.amazonaws.com'
    env.hosts = [env.master]

@task
def testing():
    """TEMPORARY - for testing the open source deploy changes"""
    _common_env()
    env.environment = 'testing'
    env.master = 'serviceinfo-testing.caktusgroup.com'
    env.hosts = [env.master]

@task
def production():
    _common_env()
    env.environment = 'production'
    # Internal hostname for our own access to the server
    # To change the domain, see conf/pillar/<envname>/env.sls
    # env.master = 'serviceinfo.rescue.org'
    env.master = env.production_master
    env.hosts = [env.master]

def get_salt_version(command):
    """Run `command` --version, pick out the part of the output that is digits and dots,
    and return it as a string.
    If the command fails, return None.
    """
    with settings(warn_only=True):
        with hide('running', 'stdout', 'stderr'):
            result = run('%s --version' % command)
            if result.succeeded:
                return re.search(r'([\d\.]+)', result).group(0)

@task
def install_salt(version, master=False, minion=False, restart=True):
    """
    Install or upgrade Salt minion and/or master if needed.

    :param version: Version string, just numbers and dots, no leading 'v'.  E.g. "2015.5.0".
      THERE IS NO DEFAULT, you must pick a version.
    :param master: If True, include master in the install.
    :param minion: If True, include minion in the install.
    :param restart: If we don't need to reinstall a salt package, restart its server anyway.
    :returns: True if any changes were made, False if nothing was done.
    """
    master_version = None
    install_master = False
    if master:
        master_version = get_salt_version("salt")
        install_master = master_version != version
        if install_master and master_version:
            # Already installed - if Ubuntu package, uninstall current version first
            # because we're going to do a git install later
            sudo("apt-get purge salt-master -y")
        if restart and not install_master:
            sudo("service salt-master restart")

    minion_version = None
    install_minion = False
    if minion:
        minion_version = get_salt_version('salt-minion')
        install_minion = minion_version != version
        if install_minion and minion_version:
            # Already installed - if Ubuntu package, uninstall current version first
            # because we're going to do a git install later
            sudo("apt-get purge salt-minion -y")
        if restart and not install_minion:
            sudo("service salt-minion restart")

    if install_master or install_minion:
        args = []
        if install_master:
            args.append('-M')
        if not install_minion:
            args.append('-N')
        args = ' '.join(args)
        # To update local install_salt.sh: wget -O install_salt.sh https://bootstrap.saltstack.com
        # then inspect it
        put(local_path="install_salt.sh", remote_path="install_salt.sh")
        sudo("sh install_salt.sh -D {args} git v{version}".format(args=args, version=version))
        return True
    return False

@task
def setup_master():
    """Provision master with salt-master."""

    # Get config onto system before the bootstrap tries to start salt
    sudo("mkdir -p /etc/salt")
    put(local_path='conf/master.conf', remote_path="/etc/salt/master", use_sudo=True)
    # If no changes, restart salt-master anyway
    install_salt(SALT_VERSION, master=True, restart=True)

@task
def sync():
    """Rysnc local states and pillar data to the master, and checkout margarita."""
    # Check for missing local secrets so that they don't get deleted
    # project.rsync_project fails if host is not set
    sudo("mkdir -p /srv")
    if not have_secrets():
        get_secrets()
    else:
        # Check for differences in the secrets files
        for environment in [env.environment]:
            remote_file = os.path.join('/srv/pillar/', environment, 'secrets.sls')
            with lcd(os.path.join(CONF_ROOT, 'pillar', environment)):
                if files.exists(remote_file):
                    get(remote_file, 'secrets.sls.remote')
                else:
                    local('touch secrets.sls.remote')
                with settings(warn_only=True):
                    result = local('diff -u secrets.sls.remote secrets.sls')
                    if result.failed and files.exists(remote_file) and not confirm(
                            red("Above changes will be made to secrets.sls. Continue?")):
                        abort("Aborted. File have been copied to secrets.sls.remote. " +
                              "Resolve conflicts, then retry.")
                    else:
                        local("rm secrets.sls.remote")
    salt_root = CONF_ROOT if CONF_ROOT.endswith('/') else CONF_ROOT + '/'
    project.rsync_project(local_dir=salt_root, remote_dir='/tmp/salt', delete=True)
    sudo('rm -rf /srv/salt /srv/pillar')
    sudo('mv /tmp/salt/* /srv/')
    sudo('rm -rf /tmp/salt/')
    execute(margarita)

def have_secrets():
    """Check if the local secret files exist for all environments."""
    found = True
    for environment in [env.environment]:
        local_file = os.path.join(CONF_ROOT, 'pillar', environment, 'secrets.sls')
        found = found and os.path.exists(local_file)
    return found

@task
def get_secrets():
    """Grab the latest secrets file from the master."""
    require('environment')
    for environment in [env.environment]:
        local_file = os.path.join(CONF_ROOT, 'pillar', environment, 'secrets.sls')
        if os.path.exists(local_file):
            local('cp {0} {0}.bak'.format(local_file))
        remote_file = os.path.join('/srv/pillar/', environment, 'secrets.sls')
        get(remote_file, local_file)

@task
def setup_minion(*roles):
    """Setup a minion server with a set of roles."""
    require('environment')
    for r in roles:
        if r not in VALID_ROLES:
            abort('%s is not a valid server role for this project.' % r)
    config = {
        'master': 'localhost' if env.master == env.host else env.master,
        'output': 'mixed',
        'grains': {
            'environment': env.environment,
            'roles': list(roles),
        },
        'mine_functions': {
            'network.interfaces': [],
            'network.ip_addrs': []
        },
    }
    _, path = tempfile.mkstemp()
    with open(path, 'w') as f:
        yaml.dump(config, f, default_flow_style=False)
    sudo("mkdir -p /etc/salt")
    put(local_path=path, remote_path="/etc/salt/minion", use_sudo=True)
    # install salt minion if it's not there already
    install_salt(SALT_VERSION, master=False, minion=True, restart=True)
    # queries server for its fully qualified domain name to get minion id
    key_name = run('python -c "import socket; print socket.getfqdn()"')
    execute(accept_key, key_name)

@task
def add_role(name):
    """Add a role to an exising minion configuration."""
    if name not in VALID_ROLES:
        abort('%s is not a valid server role for this project.' % name)
    _, path = tempfile.mkstemp()
    get("/etc/salt/minion", path)
    with open(path, 'r') as f:
        config = yaml.safe_load(f)
    grains = config.get('grains', {})
    roles = grains.get('roles', [])
    if name not in roles:
        roles.append(name)
    else:
        abort('Server is already configured with the %s role.' % name)
    grains['roles'] = roles
    config['grains'] = grains
    with open(path, 'w') as f:
        yaml.dump(config, f, default_flow_style=False)
    put(local_path=path, remote_path="/etc/salt/minion", use_sudo=True)
    sudo('service salt-minion restart')

@task
def margarita():
    require('environment')
    execute(state, 'margarita')
    sudo('service salt-master restart')

@task
def salt(cmd, target="'*'", loglevel=DEFAULT_SALT_LOGLEVEL):
    """Run arbitrary salt commands."""
    with settings(warn_only=True):
        sudo("salt {0} -l{1} {2} ".format(target, loglevel, cmd))

@task
def highstate(target="'*'", loglevel=DEFAULT_SALT_LOGLEVEL):
    """Run highstate on master."""
    print("This can take a long time without output, be patient")
    execute(salt, 'state.highstate', target, loglevel, hosts=[env.master])

@task
def state(name, target="'*'"):
    salt('state.sls {}'.format(name), target)

@task
def accept_key(name):
    """Accept minion key on master."""
    sudo('salt-key --accept={0} -y'.format(name))
    sudo('salt-key -L')

@task
def delete_key(name):
    """Delete specific key on master."""
    sudo('salt-key -L')
    sudo('salt-key --delete={0} -y'.format(name))
    sudo('salt-key -L')

@task
def deploy(loglevel=DEFAULT_SALT_LOGLEVEL):
    """Deploy to a given environment by pushing the latest states and executing the highstate."""
    require('environment')
    if env.environment != "local":
        sync()
    target = "-G 'environment:{0}'".format(env.environment)
    salt('saltutil.sync_all', target, loglevel)
    highstate(target)

@task
def manage_run(command):
    """
    Run a Django management command on the remote server.
    """
    require('environment')
    # Setup the call
    settings = '{0}.settings.{1}'.format(PROJECT_NAME, env.environment)
    manage_sh = u"DJANGO_SETTINGS_MODULE={0} /var/www/{1}/manage.sh ".format(settings, PROJECT_NAME)
    sudo(manage_sh + command, user=PROJECT_NAME)

@task
def ssh():
    """
    Convenience task to ssh to whatever host has been selected.

    E.g. ``fab production ssh``
    """
    require('environment')
    local("ssh %s" % env.hosts[0])

def _get_db_wrapper():
    return '/home/trawick/with_db.sh'

@task
def get_db_dump(clean=False):
    """Get db dump of remote enviroment."""
    require('environment')
    db_name = '%(project)s_%(environment)s' % env
    dump_file = db_name + '.sql' % env
    project_root = os.path.join('/var', 'www', env.project)
    temp_file = os.path.join(project_root, dump_file)
    flags = '-Ox'
    if clean:
        flags += 'c'
    dump_command = '%s pg_dump %s %s -U %s > %s' % (_get_db_wrapper(), flags, db_name, db_name, temp_file)
    with settings(host_string=env.master):
        sudo(dump_command, user=env.project)
        get(temp_file, dump_file)

@task
def reset_local_db():
    """ Reset local database from remote host """
    require('environment')
    question = 'Are you sure you want to reset your local ' \
               'database with the %(environment)s database?' % env
    if not confirm(question, default=False):
        abort('Local database reset aborted.')
    remote_db_name = '%(project)s_%(environment)s' % env
    db_dump_name = remote_db_name + '.sql'
    local_db_name = env.project
    get_db_dump()
    with settings(warn_only=True):
        local('dropdb %s' % local_db_name)
    local('createdb -E UTF-8 %s' % local_db_name)
    local('psql %s -c "CREATE EXTENSION postgis;"' % local_db_name)
    local('cat %s | psql %s' % (db_dump_name, local_db_name))

@task
def reset_local_media(service_info_directory):
    """ Reset local media from remote host """
    require('environment')
    media_source = os.path.join('/var', 'www', env.project, 'public', 'media')
    media_target = os.path.join(service_info_directory, 'public')
    with settings():
        local("rsync -rvaz %s:%s %s" % (env.master, media_source, media_target))

@task
def refresh_environment():
    require('environment')
    if env.environment == 'production':
        abort('Production cannot be refreshed!')

    db_name = '%(project)s_production' % env
    dump_file_name = '%s.sql' % db_name
    full_dump_file_path = os.path.join(env.project_root, dump_file_name)
    prod_host = env.production_master
    db_wrapper = _get_db_wrapper()

    with settings(host_string=prod_host):
        sudo('%s pg_dump -Ox %s -U %s > %s' % (db_wrapper, db_name, db_name, full_dump_file_path))

    sudo('supervisorctl stop all')
    db_name = db_user = '%s_%s' % (env.project, env.environment)
    media_full_path = '%s/public/media' % env.project_root
    with cd('/tmp'):
        run('scp %s:%s %s' % (prod_host, full_dump_file_path, dump_file_name))
        sudo('%s dropdb %s_backup' % (db_wrapper, db_name), user='postgres')
        sudo('psql -c "alter database %s rename to %s_backup"' % (db_name, db_name),
            user='postgres')
        sudo('%s createdb -E UTF-8 -O %s %s' % (db_wrapper, db_user, db_name), user='postgres')
        sudo('%s psql %s -c "CREATE EXTENSION postgis;"' % (db_wrapper, db_name), user='postgres')
        sudo('%s psql -U %s -d %s -f %s' % (db_wrapper, db_user, db_name, dump_file_name))
        run('rsync -zPae ssh %s:%s .' % (prod_host, media_full_path))
        sudo('rm -rf %s.backup' % media_full_path)
        sudo('mv %s %s.backup' % (media_full_path, media_full_path))
        sudo('mv media %s' % media_full_path)
        sudo('chown -R %s:%s %s' % (env.project, env.project, media_full_path))

    manage_run("migrate")
    manage_run("rebuild_index --noinput")
    sudo('supervisorctl start all')
<commit_msg>Use deployed command instead of relying on /home/trawick<commit_after>import os
import re
import tempfile

import yaml

from fabric.api import cd, env, execute, get, hide, lcd, local, put, require, run, settings, sudo, task
from fabric.colors import red
from fabric.contrib import files, project
from fabric.contrib.console import confirm
from fabric.utils import abort

DEFAULT_SALT_LOGLEVEL = 'info'
SALT_VERSION = '2015.5.1'
PROJECT_NAME = "service_info"
PROJECT_ROOT = os.path.dirname(__file__)
CONF_ROOT = os.path.join(PROJECT_ROOT, 'conf')

VALID_ROLES = (
    'salt-master',
    'web',
    'worker',
    'balancer',
    'queue',
    'cache',
    'beat',
)

def _common_env():
    env.project = PROJECT_NAME
    # refresh_environment() needs this even though the task is run against staging
    env.production_master = 'ec2-54-93-51-232.eu-central-1.compute.amazonaws.com'

@task
def staging():
    _common_env()
    env.environment = 'staging'
    # 54.93.66.254 Staging server on AWS Frankfurt
    # Internal hostname for our own access to the server
    # To change the domain, see conf/pillar/<envname>/env.sls
    # env.master = 'serviceinfo-staging.rescue.org'
    env.master = 'ec2-54-93-66-254.eu-central-1.compute.amazonaws.com'
    env.hosts = [env.master]

@task
def testing():
    """TEMPORARY - for testing the open source deploy changes"""
    _common_env()
    env.environment = 'testing'
    env.master = 'serviceinfo-testing.caktusgroup.com'
    env.hosts = [env.master]

@task
def production():
    _common_env()
    env.environment = 'production'
    # Internal hostname for our own access to the server
    # To change the domain, see conf/pillar/<envname>/env.sls
    # env.master = 'serviceinfo.rescue.org'
    env.master = env.production_master
    env.hosts = [env.master]

def get_salt_version(command):
    """Run `command` --version, pick out the part of the output that is digits and dots,
    and return it as a string.
    If the command fails, return None.
    """
    with settings(warn_only=True):
        with hide('running', 'stdout', 'stderr'):
            result = run('%s --version' % command)
            if result.succeeded:
                return re.search(r'([\d\.]+)', result).group(0)

@task
def install_salt(version, master=False, minion=False, restart=True):
    """
    Install or upgrade Salt minion and/or master if needed.

    :param version: Version string, just numbers and dots, no leading 'v'.  E.g. "2015.5.0".
      THERE IS NO DEFAULT, you must pick a version.
    :param master: If True, include master in the install.
    :param minion: If True, include minion in the install.
    :param restart: If we don't need to reinstall a salt package, restart its server anyway.
    :returns: True if any changes were made, False if nothing was done.
    """
    master_version = None
    install_master = False
    if master:
        master_version = get_salt_version("salt")
        install_master = master_version != version
        if install_master and master_version:
            # Already installed - if Ubuntu package, uninstall current version first
            # because we're going to do a git install later
            sudo("apt-get purge salt-master -y")
        if restart and not install_master:
            sudo("service salt-master restart")

    minion_version = None
    install_minion = False
    if minion:
        minion_version = get_salt_version('salt-minion')
        install_minion = minion_version != version
        if install_minion and minion_version:
            # Already installed - if Ubuntu package, uninstall current version first
            # because we're going to do a git install later
            sudo("apt-get purge salt-minion -y")
        if restart and not install_minion:
            sudo("service salt-minion restart")

    if install_master or install_minion:
        args = []
        if install_master:
            args.append('-M')
        if not install_minion:
            args.append('-N')
        args = ' '.join(args)
        # To update local install_salt.sh: wget -O install_salt.sh https://bootstrap.saltstack.com
        # then inspect it
        put(local_path="install_salt.sh", remote_path="install_salt.sh")
        sudo("sh install_salt.sh -D {args} git v{version}".format(args=args, version=version))
        return True
    return False

@task
def setup_master():
    """Provision master with salt-master."""

    # Get config onto system before the bootstrap tries to start salt
    sudo("mkdir -p /etc/salt")
    put(local_path='conf/master.conf', remote_path="/etc/salt/master", use_sudo=True)
    # If no changes, restart salt-master anyway
    install_salt(SALT_VERSION, master=True, restart=True)

@task
def sync():
    """Rysnc local states and pillar data to the master, and checkout margarita."""
    # Check for missing local secrets so that they don't get deleted
    # project.rsync_project fails if host is not set
    sudo("mkdir -p /srv")
    if not have_secrets():
        get_secrets()
    else:
        # Check for differences in the secrets files
        for environment in [env.environment]:
            remote_file = os.path.join('/srv/pillar/', environment, 'secrets.sls')
            with lcd(os.path.join(CONF_ROOT, 'pillar', environment)):
                if files.exists(remote_file):
                    get(remote_file, 'secrets.sls.remote')
                else:
                    local('touch secrets.sls.remote')
                with settings(warn_only=True):
                    result = local('diff -u secrets.sls.remote secrets.sls')
                    if result.failed and files.exists(remote_file) and not confirm(
                            red("Above changes will be made to secrets.sls. Continue?")):
                        abort("Aborted. File have been copied to secrets.sls.remote. " +
                              "Resolve conflicts, then retry.")
                    else:
                        local("rm secrets.sls.remote")
    salt_root = CONF_ROOT if CONF_ROOT.endswith('/') else CONF_ROOT + '/'
    project.rsync_project(local_dir=salt_root, remote_dir='/tmp/salt', delete=True)
    sudo('rm -rf /srv/salt /srv/pillar')
    sudo('mv /tmp/salt/* /srv/')
    sudo('rm -rf /tmp/salt/')
    execute(margarita)

def have_secrets():
    """Check if the local secret files exist for all environments."""
    found = True
    for environment in [env.environment]:
        local_file = os.path.join(CONF_ROOT, 'pillar', environment, 'secrets.sls')
        found = found and os.path.exists(local_file)
    return found

@task
def get_secrets():
    """Grab the latest secrets file from the master."""
    require('environment')
    for environment in [env.environment]:
        local_file = os.path.join(CONF_ROOT, 'pillar', environment, 'secrets.sls')
        if os.path.exists(local_file):
            local('cp {0} {0}.bak'.format(local_file))
        remote_file = os.path.join('/srv/pillar/', environment, 'secrets.sls')
        get(remote_file, local_file)

@task
def setup_minion(*roles):
    """Setup a minion server with a set of roles."""
    require('environment')
    for r in roles:
        if r not in VALID_ROLES:
            abort('%s is not a valid server role for this project.' % r)
    config = {
        'master': 'localhost' if env.master == env.host else env.master,
        'output': 'mixed',
        'grains': {
            'environment': env.environment,
            'roles': list(roles),
        },
        'mine_functions': {
            'network.interfaces': [],
            'network.ip_addrs': []
        },
    }
    _, path = tempfile.mkstemp()
    with open(path, 'w') as f:
        yaml.dump(config, f, default_flow_style=False)
    sudo("mkdir -p /etc/salt")
    put(local_path=path, remote_path="/etc/salt/minion", use_sudo=True)
    # install salt minion if it's not there already
    install_salt(SALT_VERSION, master=False, minion=True, restart=True)
    # queries server for its fully qualified domain name to get minion id
    key_name = run('python -c "import socket; print socket.getfqdn()"')
    execute(accept_key, key_name)

@task
def add_role(name):
    """Add a role to an exising minion configuration."""
    if name not in VALID_ROLES:
        abort('%s is not a valid server role for this project.' % name)
    _, path = tempfile.mkstemp()
    get("/etc/salt/minion", path)
    with open(path, 'r') as f:
        config = yaml.safe_load(f)
    grains = config.get('grains', {})
    roles = grains.get('roles', [])
    if name not in roles:
        roles.append(name)
    else:
        abort('Server is already configured with the %s role.' % name)
    grains['roles'] = roles
    config['grains'] = grains
    with open(path, 'w') as f:
        yaml.dump(config, f, default_flow_style=False)
    put(local_path=path, remote_path="/etc/salt/minion", use_sudo=True)
    sudo('service salt-minion restart')

@task
def margarita():
    require('environment')
    execute(state, 'margarita')
    sudo('service salt-master restart')

@task
def salt(cmd, target="'*'", loglevel=DEFAULT_SALT_LOGLEVEL):
    """Run arbitrary salt commands."""
    with settings(warn_only=True):
        sudo("salt {0} -l{1} {2} ".format(target, loglevel, cmd))

@task
def highstate(target="'*'", loglevel=DEFAULT_SALT_LOGLEVEL):
    """Run highstate on master."""
    print("This can take a long time without output, be patient")
    execute(salt, 'state.highstate', target, loglevel, hosts=[env.master])

@task
def state(name, target="'*'"):
    salt('state.sls {}'.format(name), target)

@task
def accept_key(name):
    """Accept minion key on master."""
    sudo('salt-key --accept={0} -y'.format(name))
    sudo('salt-key -L')

@task
def delete_key(name):
    """Delete specific key on master."""
    sudo('salt-key -L')
    sudo('salt-key --delete={0} -y'.format(name))
    sudo('salt-key -L')

@task
def deploy(loglevel=DEFAULT_SALT_LOGLEVEL):
    """Deploy to a given environment by pushing the latest states and executing the highstate."""
    require('environment')
    if env.environment != "local":
        sync()
    target = "-G 'environment:{0}'".format(env.environment)
    salt('saltutil.sync_all', target, loglevel)
    highstate(target)

@task
def manage_run(command):
    """
    Run a Django management command on the remote server.
    """
    require('environment')
    # Setup the call
    settings = '{0}.settings.{1}'.format(PROJECT_NAME, env.environment)
    manage_sh = u"DJANGO_SETTINGS_MODULE={0} /var/www/{1}/manage.sh ".format(settings, PROJECT_NAME)
    sudo(manage_sh + command, user=PROJECT_NAME)

@task
def ssh():
    """
    Convenience task to ssh to whatever host has been selected.

    E.g. ``fab production ssh``
    """
    require('environment')
    local("ssh %s" % env.hosts[0])

def _get_db_wrapper():
    return '/var/www/%s/run_with_db.sh' % env.project

@task
def get_db_dump(clean=False):
    """Get db dump of remote enviroment."""
    require('environment')
    db_name = '%(project)s_%(environment)s' % env
    dump_file = db_name + '.sql' % env
    project_root = os.path.join('/var', 'www', env.project)
    temp_file = os.path.join(project_root, dump_file)
    flags = '-Ox'
    if clean:
        flags += 'c'
    dump_command = '%s pg_dump %s %s -U %s > %s' % (_get_db_wrapper(), flags, db_name, db_name, temp_file)
    with settings(host_string=env.master):
        sudo(dump_command, user=env.project)
        get(temp_file, dump_file)

@task
def reset_local_db():
    """ Reset local database from remote host """
    require('environment')
    question = 'Are you sure you want to reset your local ' \
               'database with the %(environment)s database?' % env
    if not confirm(question, default=False):
        abort('Local database reset aborted.')
    remote_db_name = '%(project)s_%(environment)s' % env
    db_dump_name = remote_db_name + '.sql'
    local_db_name = env.project
    get_db_dump()
    with settings(warn_only=True):
        local('dropdb %s' % local_db_name)
    local('createdb -E UTF-8 %s' % local_db_name)
    local('psql %s -c "CREATE EXTENSION postgis;"' % local_db_name)
    local('cat %s | psql %s' % (db_dump_name, local_db_name))

@task
def reset_local_media(service_info_directory):
    """ Reset local media from remote host """
    require('environment')
    media_source = os.path.join('/var', 'www', env.project, 'public', 'media')
    media_target = os.path.join(service_info_directory, 'public')
    with settings():
        local("rsync -rvaz %s:%s %s" % (env.master, media_source, media_target))

@task
def refresh_environment():
    require('environment')
    if env.environment == 'production':
        abort('Production cannot be refreshed!')

    db_name = '%(project)s_production' % env
    dump_file_name = '%s.sql' % db_name
    full_dump_file_path = os.path.join(env.project_root, dump_file_name)
    prod_host = env.production_master
    db_wrapper = _get_db_wrapper()

    with settings(host_string=prod_host):
        sudo('%s pg_dump -Ox %s -U %s > %s' % (db_wrapper, db_name, db_name, full_dump_file_path))

    sudo('supervisorctl stop all')
    db_name = db_user = '%s_%s' % (env.project, env.environment)
    media_full_path = '%s/public/media' % env.project_root
    with cd('/tmp'):
        run('scp %s:%s %s' % (prod_host, full_dump_file_path, dump_file_name))
        sudo('%s dropdb %s_backup' % (db_wrapper, db_name), user='postgres')
        sudo('psql -c "alter database %s rename to %s_backup"' % (db_name, db_name),
            user='postgres')
        sudo('%s createdb -E UTF-8 -O %s %s' % (db_wrapper, db_user, db_name), user='postgres')
        sudo('%s psql %s -c "CREATE EXTENSION postgis;"' % (db_wrapper, db_name), user='postgres')
        sudo('%s psql -U %s -d %s -f %s' % (db_wrapper, db_user, db_name, dump_file_name))
        run('rsync -zPae ssh %s:%s .' % (prod_host, media_full_path))
        sudo('rm -rf %s.backup' % media_full_path)
        sudo('mv %s %s.backup' % (media_full_path, media_full_path))
        sudo('mv media %s' % media_full_path)
        sudo('chown -R %s:%s %s' % (env.project, env.project, media_full_path))

    manage_run("migrate")
    manage_run("rebuild_index --noinput")
    sudo('supervisorctl start all')
<|endoftext|>


<commit_before># Under MIT License, see LICENSE.txt

from collections import namedtuple

class AICommand(namedtuple('AICommand', 'target kick_type kick_force charge_kick dribbler_active '
                                        'cruise_speed end_speed ball_collision pathfinder_on')):
    __slots__ = ()

class CmdBuilder:

    def __init__(self):
        # Does are the defaults values of AICommand
        # A robot should stay idle with those value
        self.target = None
        self.kick_type = None
        self.kick_force = 0
        self.charge_kick = False
        self.dribbler_active = False
        self.cruise_speed = 0
        self.end_speed = 0
        self.ball_collision = True
        self.pathfinder_on = True

    def withMoveTo(self, target, cruise_speed=1, end_speed=0, ball_collision=True):
        self.target = target
        self.cruise_speed = cruise_speed
        self.end_speed = end_speed
        self.ball_collision = ball_collision
        self.pathfinder_on = True
        return self

    def withKick(self, kick_force=1):
        self.kick_force = kick_force
        self.kick_type = 1 # For the moment we only have one type
        return self

    def withForceDribbler(self):
        self.dribbler_active = True
        return self

    def withChargeKicker(self):
        self.charge_kick = True
        return self

    def build(self):
        return AICommand(self.target,
                         self.kick_type,
                         self.kick_force,
                         self.charge_kick,
                         self.dribbler_active,
                         self.cruise_speed,
                         self.end_speed,
                         self.ball_collision,
                         self.pathfinder_on)

Idle = CmdBuilder().build()<commit_msg>minor fixes<commit_after># Under MIT License, see LICENSE.txt

from collections import namedtuple

AICommand = namedtuple('AICommand', 'target,'
                                    'kick_type,'
                                    'kick_force,'
                                    'charge_kick,'
                                    'dribbler_active,'
                                    'cruise_speed,'
                                    'end_speed,'
                                    'ball_collision,'
                                    'pathfinder_on')

class CmdBuilder:

    def __init__(self):
        # Does are the defaults values of AICommand
        # A robot should stay idle with those value
        self.target = None
        self.kick_type = None
        self.kick_force = 0
        self.charge_kick = False
        self.dribbler_active = False
        self.cruise_speed = 0
        self.end_speed = 0
        self.ball_collision = True
        self.pathfinder_on = True

    def withMoveTo(self, target, cruise_speed=1, end_speed=0, ball_collision=True):
        self.target = target
        self.cruise_speed = cruise_speed
        self.end_speed = end_speed
        self.ball_collision = ball_collision
        self.pathfinder_on = True
        return self

    def withKick(self, kick_force=1):
        self.kick_force = kick_force
        self.kick_type = 1  # For the moment we only have one type
        return self

    def withForceDribbler(self):
        self.dribbler_active = True
        return self

    def withChargeKicker(self):
        self.charge_kick = True
        return self

    def build(self):
        return AICommand(self.target,
                         self.kick_type,
                         self.kick_force,
                         self.charge_kick,
                         self.dribbler_active,
                         self.cruise_speed,
                         self.end_speed,
                         self.ball_collision,
                         self.pathfinder_on)

Idle = CmdBuilder().build()
<|endoftext|>


<commit_before>import functools
import os

from datetime import datetime
from fabric.api import local, env, lcd
from fabric.colors import yellow, red, green
from fabric.contrib import console

PROJECT_ROOT = os.path.dirname(os.path.realpath(__file__))
here = lambda *x: os.path.join(PROJECT_ROOT, *x)

def production():
    """Connection details for the ``production`` app"""
    env.slug = 'production'
    env.url = 'http://us-ignite.herokuapp.com/'
    env.app = 'us-ignite'

def is_vm():
    """Determines if the script is running in a VM"""
    return os.environ['USER'] == 'vagrant'

IS_VM = is_vm()

def only_outside_vm(function):
    """Decorator to allow only ``inside`` VM commands """
    @functools.wraps(function)
    def inner(*args, **kwargs):
        if IS_VM:
            print red('Command can only be OUTSIDE the VM.')
            return exit(1)
        return function(*args, **kwargs)
    return inner

def only_inside_vm(function):
    """Decorator to allow only ``outside`` VM commands"""
    @functools.wraps(function)
    def inner(*args, **kwargs):
        if not IS_VM:
            print red('Command can only be INSIDE the VM.')
            return exit(1)
        return function(*args, **kwargs)
    return inner

def dj_heroku(command, slug, capture=False):
    """Runs a given django management command in the given Heroku's app."""
    new_cmd = ('heroku run django-admin.py %s --settings=us_ignite.'
               'settings.%s --remote %s' % (command, slug, slug))
    return local(new_cmd, capture)

def run_heroku(cmd, slug, capture=True):
    """Runs a Heroku command with the given ``remote``"""
    return local('heroku %s --app %s' % (cmd, slug), capture=capture)

@only_outside_vm
def syncdb():
    print yellow('Syncing %s database.' % env.slug)
    dj_heroku('syncdb --noinput', env.slug)

@only_outside_vm
def shell():
    """Open a shell in the given environment."""
    dj_heroku('shell', env.slug)

def production_confirmation(function):
    """Production confirmation."""
    @functools.wraps(function)
    def wrapper(confirmation='', *args, **kwargs):
        confirmation = False if confirmation == 'False' else True
        SLUG = env.slug.upper()
        if env.slug not in ['production']:
            print red('Invaid destination: %s.' % SLUG)
            exit(3)
        if confirmation and env.slug == 'production':
            msg = red('You are about to DEPLOY %s to Heroku. Procceed?' % SLUG)
            if not console.confirm(msg):
                print yellow('Phew, aborted.')
                exit(2)
        return function(confirmation, *args, **kwargs)
    return wrapper

def _validate_pushed_commits():
    with lcd(PROJECT_ROOT):
        result = local('git log --pretty=format:"%h %s" origin/master..HEAD',
                       capture=True)
        if not result:
            return True
        print red('FAILURE: There are unpushed commits to origin/master.')
        print yellow('Make sure these commits are pushed before '
                     'deploying to heroku:')
        print yellow(result)
        exit(4)

@only_outside_vm
@production_confirmation
def deploy(confirmation):
    """Deploys the given build."""
    SLUG = env.slug.upper()
    print yellow('Deploying to %s. Because you said so.' % SLUG)
    with lcd(PROJECT_ROOT):
        print yellow('Pushing changes to %s in Heroku.' % SLUG)
        _validate_pushed_commits()
        local('git push %s master' % env.slug)
        syncdb()
    print yellow('URL: %s' % env.url)
    print red('Done! - %s' % datetime.now())

@only_inside_vm
def test(apps=''):
    """Runs the test suite."""
    local('django-admin.py test %s '
          '--settings=us_ignite.settings.testing' % apps)
<commit_msg>Add support for reseting local database.<commit_after>import functools
import os

from datetime import datetime
from fabric.api import local, env, lcd
from fabric.colors import yellow, red, green
from fabric.contrib import console

PROJECT_ROOT = os.path.dirname(os.path.realpath(__file__))
here = lambda *x: os.path.join(PROJECT_ROOT, *x)

DB_STRING = 'us_ignite'

def production():
    """Connection details for the ``production`` app"""
    env.slug = 'production'
    env.url = 'http://us-ignite.herokuapp.com/'
    env.app = 'us-ignite'

def is_vm():
    """Determines if the script is running in a VM"""
    return os.environ['USER'] == 'vagrant'

IS_VM = is_vm()

def only_outside_vm(function):
    """Decorator to allow only ``inside`` VM commands """
    @functools.wraps(function)
    def inner(*args, **kwargs):
        if IS_VM:
            print red('Command can only be OUTSIDE the VM.')
            return exit(1)
        return function(*args, **kwargs)
    return inner

def only_inside_vm(function):
    """Decorator to allow only ``outside`` VM commands"""
    @functools.wraps(function)
    def inner(*args, **kwargs):
        if not IS_VM:
            print red('Command can only be INSIDE the VM.')
            return exit(1)
        return function(*args, **kwargs)
    return inner

def dj_heroku(command, slug, capture=False):
    """Runs a given django management command in the given Heroku's app."""
    new_cmd = ('heroku run django-admin.py %s --settings=us_ignite.'
               'settings.%s --remote %s' % (command, slug, slug))
    return local(new_cmd, capture)

def run_heroku(cmd, slug, capture=True):
    """Runs a Heroku command with the given ``remote``"""
    return local('heroku %s --app %s' % (cmd, slug), capture=capture)

@only_outside_vm
def syncdb():
    print yellow('Syncing %s database.' % env.slug)
    dj_heroku('syncdb --noinput', env.slug)

@only_outside_vm
def shell():
    """Open a shell in the given environment."""
    dj_heroku('shell', env.slug)

def production_confirmation(function):
    """Production confirmation."""
    @functools.wraps(function)
    def wrapper(confirmation='', *args, **kwargs):
        confirmation = False if confirmation == 'False' else True
        SLUG = env.slug.upper()
        if env.slug not in ['production']:
            print red('Invaid destination: %s.' % SLUG)
            exit(3)
        if confirmation and env.slug == 'production':
            msg = red('You are about to DEPLOY %s to Heroku. Procceed?' % SLUG)
            if not console.confirm(msg):
                print yellow('Phew, aborted.')
                exit(2)
        return function(confirmation, *args, **kwargs)
    return wrapper

def _validate_pushed_commits():
    with lcd(PROJECT_ROOT):
        result = local('git log --pretty=format:"%h %s" origin/master..HEAD',
                       capture=True)
        if not result:
            return True
        print red('FAILURE: There are unpushed commits to origin/master.')
        print yellow('Make sure these commits are pushed before '
                     'deploying to heroku:')
        print yellow(result)
        exit(4)

@only_outside_vm
@production_confirmation
def deploy(confirmation):
    """Deploys the given build."""
    SLUG = env.slug.upper()
    print yellow('Deploying to %s. Because you said so.' % SLUG)
    with lcd(PROJECT_ROOT):
        print yellow('Pushing changes to %s in Heroku.' % SLUG)
        _validate_pushed_commits()
        local('git push %s master' % env.slug)
        syncdb()
    print yellow('URL: %s' % env.url)
    print red('Done! - %s' % datetime.now())

@only_inside_vm
def test(apps=''):
    """Runs the test suite."""
    local('django-admin.py test %s '
          '--settings=us_ignite.settings.testing' % apps)

@only_inside_vm
def drop_local_db():
    """Removes the local development database."""
    print yellow('Droping and creating a new DB.')
    local('PGPASSWORD=%(db)s dropdb %(db)s -U %(db)s -h localhost'
          % {'db': DB_STRING})
    local("PGPASSWORD=%(db)s createdb %(db)s -T template0 -E UTF-8 "
          "-l en_US.UTF-8 -O %(db)s -U %(db)s -h localhost"
          % {'db': DB_STRING})

@only_inside_vm
def reset_local_db():
    """Resets the local development DB
    Check the database has the right UTF8 encoding before running this command.
    sudo -u postgres psql --listsudo -u postgres psql --list
    """
    confirmation = red('You are about to IRREVERSIBLY clear the existing'
                       ' local database. Procceed?')
    if console.confirm(confirmation):
        drop_local_db()
        local('django-admin.py syncdb --noinput '
              '--settings=%s.settings.local' % DB_STRING)
        # local('django-admin.py migrate --noinput '
        #       '--settings=%s.settings.local' % DB_STRING)
    else:
        print yellow('Phew, aborted.')
        exit(1)
<|endoftext|>


<commit_before># Copyright 2020 Google LLC. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for the cloud docker containerization module."""

import docker
import mock
import os
import tarfile
import unittest

from tensorflow_cloud.core import containerize
from tensorflow_cloud.core import machine_config

try:
    from tensorflow.python.framework.versions import VERSION
except ImportError:
    # Use the latest TF docker image if a local installation is not available.
    VERSION = "latest"

class TestContainerize(unittest.TestCase):
    def setup(self):
        self.entry_point = "sample.py"
        self.chief_config = machine_config.COMMON_MACHINE_CONFIGS["K80_1X"]
        self.worker_config = machine_config.COMMON_MACHINE_CONFIGS["K80_1X"]
        self.entry_point_dir = "."
        self.mock_registry = "gcr.io/my-project"
        self.project_id = "my-project"

    def cleanup(self, docker_file):
        os.remove(docker_file)

    def assert_docker_file(self, expected_lines, docker_file):
        with open(docker_file, "r") as f:
            actual_lines = f.readlines()
            self.assertListEqual(expected_lines, actual_lines)

    def test_create_docker_file_defaults(self):
        self.setup()
        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
        )
        lcb._create_docker_file()

        expected_docker_file_lines = [
            "FROM tensorflow/tensorflow:{}-gpu\n".format(VERSION),
            "WORKDIR /app/\n",
            "COPY /app/ /app/\n",
            'ENTRYPOINT ["python", "sample.py"]',
        ]
        self.assert_docker_file(expected_docker_file_lines, lcb.docker_file_path)
        self.cleanup(lcb.docker_file_path)

    def test_create_docker_with_requirements(self):
        self.setup()
        req_file = "requirements.txt"
        with open(req_file, "w") as f:
            f.writelines(["tensorflow-datasets"])

        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
            requirements_txt=req_file,
        )
        lcb._create_docker_file()

        expected_docker_file_lines = [
            "FROM tensorflow/tensorflow:{}-gpu\n".format(VERSION),
            "WORKDIR /app/\n",
            "COPY /app/requirements.txt /app/requirements.txt\n",
            "RUN if [ -e requirements.txt ]; "
            "then pip install --no-cache -r requirements.txt; fi\n",
            "COPY /app/ /app/\n",
            'ENTRYPOINT ["python", "sample.py"]',
        ]
        self.assert_docker_file(expected_docker_file_lines, lcb.docker_file_path)

        os.remove(req_file)
        self.cleanup(lcb.docker_file_path)

    def test_create_docker_file_with_destination_dir(self):
        self.setup()
        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
            destination_dir="/my_app/temp/",
        )
        lcb._create_docker_file()

        expected_docker_file_lines = [
            "FROM tensorflow/tensorflow:{}-gpu\n".format(VERSION),
            "WORKDIR /my_app/temp/\n",
            "COPY /my_app/temp/ /my_app/temp/\n",
            'ENTRYPOINT ["python", "sample.py"]',
        ]
        self.assert_docker_file(expected_docker_file_lines, lcb.docker_file_path)
        self.cleanup(lcb.docker_file_path)

    def test_check_docker_base_image_nightly(self):
        self.setup()
        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
            docker_base_image="tensorflow/tensorflow:2.3.0-dev20200605",
        )
        # Verify the docker base image is identified as nonexistent
        self.assertTrue(not lcb._base_image_exist())

        # Verify that the dockerfile fetches the latest tfnightly
        lcb._create_docker_file()
        expected_docker_file_lines = [
            "FROM tensorflow/tensorflow:nightly\n",
            "WORKDIR /app/\n",
            "COPY /app/ /app/\n",
            'ENTRYPOINT ["python", "sample.py"]',
        ]
        self.assert_docker_file(expected_docker_file_lines, lcb.docker_file_path)
        self.cleanup(lcb.docker_file_path)

    def test_check_nonexistent_docker_image(self):
        self.setup()
        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
            docker_base_image="tensorflow/tensorflow:21",
        )
        # Verify the docker base image is identified as nonexistent
        self.assertTrue(not lcb._base_image_exist())

        # Verify that the dockerfile fetches the latest stable image
        lcb._create_docker_file()
        expected_docker_file_lines = [
            "FROM tensorflow/tensorflow:latest\n",
            "WORKDIR /app/\n",
            "COPY /app/ /app/\n",
            'ENTRYPOINT ["python", "sample.py"]',
        ]
        self.assert_docker_file(expected_docker_file_lines, lcb.docker_file_path)
        self.cleanup(lcb.docker_file_path)

    def test_create_docker_file_with_docker_base_image(self):
        self.setup()
        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
            docker_base_image="tensorflow/tensorflow:latest",
        )
        lcb._create_docker_file()

        expected_docker_file_lines = [
            "FROM tensorflow/tensorflow:latest\n",
            "WORKDIR /app/\n",
            "COPY /app/ /app/\n",
            'ENTRYPOINT ["python", "sample.py"]',
        ]
        self.assert_docker_file(expected_docker_file_lines, lcb.docker_file_path)
        self.cleanup(lcb.docker_file_path)

    def test_create_docker_file_with_cpu_config(self):
        self.setup()
        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            machine_config.COMMON_MACHINE_CONFIGS["CPU"],
            self.worker_config,
            self.mock_registry,
            self.project_id,
        )
        lcb._create_docker_file()

        expected_docker_file_lines = [
            "FROM tensorflow/tensorflow:{}\n".format(VERSION),
            "WORKDIR /app/\n",
            "COPY /app/ /app/\n",
            'ENTRYPOINT ["python", "sample.py"]',
        ]
        self.assert_docker_file(expected_docker_file_lines, lcb.docker_file_path)
        self.cleanup(lcb.docker_file_path)

    def test_create_docker_file_with_tpu_config(self):
        self.setup()
        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            machine_config.COMMON_MACHINE_CONFIGS["CPU"],
            machine_config.COMMON_MACHINE_CONFIGS["TPU"],
            self.mock_registry,
            self.project_id,
        )
        lcb._create_docker_file()

        expected_docker_file_lines = [
            "FROM tensorflow/tensorflow:{}\n".format(VERSION),
            "WORKDIR /app/\n",
            "RUN pip install cloud-tpu-client\n",
            "COPY /app/ /app/\n",
            'ENTRYPOINT ["python", "sample.py"]',
        ]
        self.assert_docker_file(expected_docker_file_lines, lcb.docker_file_path)
        self.cleanup(lcb.docker_file_path)

    def test_get_file_path_map_defaults(self):
        self.setup()
        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
        )
        lcb._create_docker_file()
        file_map = lcb._get_file_path_map()

        self.assertDictEqual(
            file_map,
            {lcb.docker_file_path: "Dockerfile", self.entry_point_dir: "/app/"},
        )

        self.cleanup(lcb.docker_file_path)

    def test_get_file_path_map_with_requirements(self):
        self.setup()
        req_file = "requirements.txt"
        with open(req_file, "w") as f:
            f.writelines(["tensorflow-datasets"])

        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
            requirements_txt=req_file,
        )
        lcb._create_docker_file()
        file_map = lcb._get_file_path_map()

        self.assertDictEqual(
            file_map,
            {
                lcb.docker_file_path: "Dockerfile",
                req_file: "/app/requirements.txt",
                self.entry_point_dir: "/app/",
            },
        )

        os.remove(req_file)
        self.cleanup(lcb.docker_file_path)

    def test_get_file_path_map_with_destination_dir(self):
        self.setup()

        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
            destination_dir="/my_app/temp/",
        )
        lcb._create_docker_file()
        file_map = lcb._get_file_path_map()

        self.assertDictEqual(
            file_map,
            {lcb.docker_file_path: "Dockerfile", self.entry_point_dir: "/my_app/temp/"},
        )

        self.cleanup(lcb.docker_file_path)

    def test_get_file_path_map_with_wrapped_entry_point(self):
        self.setup()

        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            self.entry_point,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
            destination_dir="/my_app/temp/",
        )
        lcb._create_docker_file()
        file_map = lcb._get_file_path_map()

        self.assertDictEqual(
            file_map,
            {
                lcb.docker_file_path: "Dockerfile",
                self.entry_point_dir: "/my_app/temp/",
                self.entry_point: "/my_app/temp/sample.py",
            },
        )

        self.cleanup(lcb.docker_file_path)

    @mock.patch("tensorflow_cloud.core.containerize.logger")
    @mock.patch("docker.APIClient")
    def test_get_docker_image(self, MockAPIClient, MockLogger):
        self.setup()
        mock_registry = "gcr.io/my-project"
        mock_img_tag = mock_registry + "/tensorflow-train:abcde"

        # Verify mocking is correct and mock img tag.
        assert MockAPIClient is docker.APIClient
        assert MockLogger is containerize.logger
        docker_client = MockAPIClient.return_value

        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            self.entry_point,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
            destination_dir="/my_app/temp/",
        )

        def _mock_generate_name():
            return mock_img_tag

        lcb._generate_name = _mock_generate_name
        lcb._get_file_path_map = mock.Mock()
        lcb._get_file_path_map.return_value = {}
        lcb.tar_file_path = ""
        with mock.patch("builtins.open", mock.mock_open(read_data="")) as mock_file:
            img_tag = lcb.get_docker_image()

        self.assertEqual(img_tag, mock_img_tag)

        # Verify docker APIClient is invoked as expected.
        self.assertEqual(MockAPIClient.call_count, 1)
        _, kwargs = MockAPIClient.call_args
        self.assertDictEqual(kwargs, {"version": "auto"})

        # Verify APIClient().build is invoked as expected.
        self.assertEqual(docker_client.build.call_count, 1)
        _, kwargs = docker_client.build.call_args
        expected = {
            "path": ".",
            "custom_context": True,
            "encoding": "utf-8",
            "tag": img_tag,
        }
        self.assertTrue(set(expected.items()).issubset(set(kwargs.items())))

        # Verify APIClient().push is invoked as expected.
        self.assertEqual(docker_client.push.call_count, 1)
        args, kwargs = docker_client.push.call_args
        self.assertListEqual(list(args), [img_tag])
        self.assertDictEqual(kwargs, {"decode": True, "stream": True})

        # Verify logger info calls.
        self.assertEqual(MockLogger.info.call_count, 2)
        MockLogger.info.assert_has_calls(
            [
                mock.call(r"Building docker image: " + img_tag),
                mock.call(r"Publishing docker image: " + img_tag),
            ]
        )
        self.cleanup(lcb.docker_file_path)

if __name__ == "__main__":
    unittest.main()
<commit_msg>Fix remaining TF version import<commit_after># Copyright 2020 Google LLC. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for the cloud docker containerization module."""

import docker
import mock
import os
import tarfile
import unittest

from tensorflow_cloud.core import containerize
from tensorflow_cloud.core import machine_config

try:
    from tensorflow import __version__ as VERSION
except ImportError:
    # Use the latest TF docker image if a local installation is not available.
    VERSION = "latest"

class TestContainerize(unittest.TestCase):
    def setup(self):
        self.entry_point = "sample.py"
        self.chief_config = machine_config.COMMON_MACHINE_CONFIGS["K80_1X"]
        self.worker_config = machine_config.COMMON_MACHINE_CONFIGS["K80_1X"]
        self.entry_point_dir = "."
        self.mock_registry = "gcr.io/my-project"
        self.project_id = "my-project"

    def cleanup(self, docker_file):
        os.remove(docker_file)

    def assert_docker_file(self, expected_lines, docker_file):
        with open(docker_file, "r") as f:
            actual_lines = f.readlines()
            self.assertListEqual(expected_lines, actual_lines)

    def test_create_docker_file_defaults(self):
        self.setup()
        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
        )
        lcb._create_docker_file()

        expected_docker_file_lines = [
            "FROM tensorflow/tensorflow:{}-gpu\n".format(VERSION),
            "WORKDIR /app/\n",
            "COPY /app/ /app/\n",
            'ENTRYPOINT ["python", "sample.py"]',
        ]
        self.assert_docker_file(expected_docker_file_lines, lcb.docker_file_path)
        self.cleanup(lcb.docker_file_path)

    def test_create_docker_with_requirements(self):
        self.setup()
        req_file = "requirements.txt"
        with open(req_file, "w") as f:
            f.writelines(["tensorflow-datasets"])

        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
            requirements_txt=req_file,
        )
        lcb._create_docker_file()

        expected_docker_file_lines = [
            "FROM tensorflow/tensorflow:{}-gpu\n".format(VERSION),
            "WORKDIR /app/\n",
            "COPY /app/requirements.txt /app/requirements.txt\n",
            "RUN if [ -e requirements.txt ]; "
            "then pip install --no-cache -r requirements.txt; fi\n",
            "COPY /app/ /app/\n",
            'ENTRYPOINT ["python", "sample.py"]',
        ]
        self.assert_docker_file(expected_docker_file_lines, lcb.docker_file_path)

        os.remove(req_file)
        self.cleanup(lcb.docker_file_path)

    def test_create_docker_file_with_destination_dir(self):
        self.setup()
        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
            destination_dir="/my_app/temp/",
        )
        lcb._create_docker_file()

        expected_docker_file_lines = [
            "FROM tensorflow/tensorflow:{}-gpu\n".format(VERSION),
            "WORKDIR /my_app/temp/\n",
            "COPY /my_app/temp/ /my_app/temp/\n",
            'ENTRYPOINT ["python", "sample.py"]',
        ]
        self.assert_docker_file(expected_docker_file_lines, lcb.docker_file_path)
        self.cleanup(lcb.docker_file_path)

    def test_check_docker_base_image_nightly(self):
        self.setup()
        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
            docker_base_image="tensorflow/tensorflow:2.3.0-dev20200605",
        )
        # Verify the docker base image is identified as nonexistent
        self.assertTrue(not lcb._base_image_exist())

        # Verify that the dockerfile fetches the latest tfnightly
        lcb._create_docker_file()
        expected_docker_file_lines = [
            "FROM tensorflow/tensorflow:nightly\n",
            "WORKDIR /app/\n",
            "COPY /app/ /app/\n",
            'ENTRYPOINT ["python", "sample.py"]',
        ]
        self.assert_docker_file(expected_docker_file_lines, lcb.docker_file_path)
        self.cleanup(lcb.docker_file_path)

    def test_check_nonexistent_docker_image(self):
        self.setup()
        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
            docker_base_image="tensorflow/tensorflow:21",
        )
        # Verify the docker base image is identified as nonexistent
        self.assertTrue(not lcb._base_image_exist())

        # Verify that the dockerfile fetches the latest stable image
        lcb._create_docker_file()
        expected_docker_file_lines = [
            "FROM tensorflow/tensorflow:latest\n",
            "WORKDIR /app/\n",
            "COPY /app/ /app/\n",
            'ENTRYPOINT ["python", "sample.py"]',
        ]
        self.assert_docker_file(expected_docker_file_lines, lcb.docker_file_path)
        self.cleanup(lcb.docker_file_path)

    def test_create_docker_file_with_docker_base_image(self):
        self.setup()
        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
            docker_base_image="tensorflow/tensorflow:latest",
        )
        lcb._create_docker_file()

        expected_docker_file_lines = [
            "FROM tensorflow/tensorflow:latest\n",
            "WORKDIR /app/\n",
            "COPY /app/ /app/\n",
            'ENTRYPOINT ["python", "sample.py"]',
        ]
        self.assert_docker_file(expected_docker_file_lines, lcb.docker_file_path)
        self.cleanup(lcb.docker_file_path)

    def test_create_docker_file_with_cpu_config(self):
        self.setup()
        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            machine_config.COMMON_MACHINE_CONFIGS["CPU"],
            self.worker_config,
            self.mock_registry,
            self.project_id,
        )
        lcb._create_docker_file()

        expected_docker_file_lines = [
            "FROM tensorflow/tensorflow:{}\n".format(VERSION),
            "WORKDIR /app/\n",
            "COPY /app/ /app/\n",
            'ENTRYPOINT ["python", "sample.py"]',
        ]
        self.assert_docker_file(expected_docker_file_lines, lcb.docker_file_path)
        self.cleanup(lcb.docker_file_path)

    def test_create_docker_file_with_tpu_config(self):
        self.setup()
        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            machine_config.COMMON_MACHINE_CONFIGS["CPU"],
            machine_config.COMMON_MACHINE_CONFIGS["TPU"],
            self.mock_registry,
            self.project_id,
        )
        lcb._create_docker_file()

        expected_docker_file_lines = [
            "FROM tensorflow/tensorflow:{}\n".format(VERSION),
            "WORKDIR /app/\n",
            "RUN pip install cloud-tpu-client\n",
            "COPY /app/ /app/\n",
            'ENTRYPOINT ["python", "sample.py"]',
        ]
        self.assert_docker_file(expected_docker_file_lines, lcb.docker_file_path)
        self.cleanup(lcb.docker_file_path)

    def test_get_file_path_map_defaults(self):
        self.setup()
        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
        )
        lcb._create_docker_file()
        file_map = lcb._get_file_path_map()

        self.assertDictEqual(
            file_map,
            {lcb.docker_file_path: "Dockerfile", self.entry_point_dir: "/app/"},
        )

        self.cleanup(lcb.docker_file_path)

    def test_get_file_path_map_with_requirements(self):
        self.setup()
        req_file = "requirements.txt"
        with open(req_file, "w") as f:
            f.writelines(["tensorflow-datasets"])

        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
            requirements_txt=req_file,
        )
        lcb._create_docker_file()
        file_map = lcb._get_file_path_map()

        self.assertDictEqual(
            file_map,
            {
                lcb.docker_file_path: "Dockerfile",
                req_file: "/app/requirements.txt",
                self.entry_point_dir: "/app/",
            },
        )

        os.remove(req_file)
        self.cleanup(lcb.docker_file_path)

    def test_get_file_path_map_with_destination_dir(self):
        self.setup()

        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            None,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
            destination_dir="/my_app/temp/",
        )
        lcb._create_docker_file()
        file_map = lcb._get_file_path_map()

        self.assertDictEqual(
            file_map,
            {lcb.docker_file_path: "Dockerfile", self.entry_point_dir: "/my_app/temp/"},
        )

        self.cleanup(lcb.docker_file_path)

    def test_get_file_path_map_with_wrapped_entry_point(self):
        self.setup()

        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            self.entry_point,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
            destination_dir="/my_app/temp/",
        )
        lcb._create_docker_file()
        file_map = lcb._get_file_path_map()

        self.assertDictEqual(
            file_map,
            {
                lcb.docker_file_path: "Dockerfile",
                self.entry_point_dir: "/my_app/temp/",
                self.entry_point: "/my_app/temp/sample.py",
            },
        )

        self.cleanup(lcb.docker_file_path)

    @mock.patch("tensorflow_cloud.core.containerize.logger")
    @mock.patch("docker.APIClient")
    def test_get_docker_image(self, MockAPIClient, MockLogger):
        self.setup()
        mock_registry = "gcr.io/my-project"
        mock_img_tag = mock_registry + "/tensorflow-train:abcde"

        # Verify mocking is correct and mock img tag.
        assert MockAPIClient is docker.APIClient
        assert MockLogger is containerize.logger
        docker_client = MockAPIClient.return_value

        lcb = containerize.LocalContainerBuilder(
            self.entry_point,
            self.entry_point,
            self.chief_config,
            self.worker_config,
            self.mock_registry,
            self.project_id,
            destination_dir="/my_app/temp/",
        )

        def _mock_generate_name():
            return mock_img_tag

        lcb._generate_name = _mock_generate_name
        lcb._get_file_path_map = mock.Mock()
        lcb._get_file_path_map.return_value = {}
        lcb.tar_file_path = ""
        with mock.patch("builtins.open", mock.mock_open(read_data="")) as mock_file:
            img_tag = lcb.get_docker_image()

        self.assertEqual(img_tag, mock_img_tag)

        # Verify docker APIClient is invoked as expected.
        self.assertEqual(MockAPIClient.call_count, 1)
        _, kwargs = MockAPIClient.call_args
        self.assertDictEqual(kwargs, {"version": "auto"})

        # Verify APIClient().build is invoked as expected.
        self.assertEqual(docker_client.build.call_count, 1)
        _, kwargs = docker_client.build.call_args
        expected = {
            "path": ".",
            "custom_context": True,
            "encoding": "utf-8",
            "tag": img_tag,
        }
        self.assertTrue(set(expected.items()).issubset(set(kwargs.items())))

        # Verify APIClient().push is invoked as expected.
        self.assertEqual(docker_client.push.call_count, 1)
        args, kwargs = docker_client.push.call_args
        self.assertListEqual(list(args), [img_tag])
        self.assertDictEqual(kwargs, {"decode": True, "stream": True})

        # Verify logger info calls.
        self.assertEqual(MockLogger.info.call_count, 2)
        MockLogger.info.assert_has_calls(
            [
                mock.call(r"Building docker image: " + img_tag),
                mock.call(r"Publishing docker image: " + img_tag),
            ]
        )
        self.cleanup(lcb.docker_file_path)

if __name__ == "__main__":
    unittest.main()
<|endoftext|>


<commit_before># Copyright 2013-2014 Massachusetts Open Cloud Contributors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the
# License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an "AS
# IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
# express or implied. See the License for the specific language
# governing permissions and limitations under the License.
"""A switch driver for the Dell Powerconnect
See the documentation for the haas.drivers package for a description of this
module's interface.
Currently the driver uses telnet to connect to the switch's console; in
the long term we want to be using SNMP.
"""
import os
import pexpect
import re
from haas.config import cfg
from haas.dev_support import no_dry_run
from haas.model import Model, Session
from sqlalchemy import *
class Nexus_Vlan(Model):
    """A VLAN for the Cisco Nexus switch
       This is used to track which vlan numbers are available; when a Network is
       created, it must allocate a Vlan, to ensure that:
       1. The VLAN number it is using is unique, and
       2. The VLAN number is actually allocated to the HaaS; on some deployments we
       may have specific vlan numbers that we are allowed to use.
       This code currently only works if there is only one Nexus switch. This is       because the Nexus switch does not support switch stacking.
    """
    vlan_no = Column(Integer, nullable=False, unique=True)
    available = Column(Boolean, nullable=False)
    def __init__(self, vlan_no):
        self.vlan_no = vlan_no
        self.available = True
        # XXX: This is pretty gross; it arguably doesn't even make sense for
        # Vlan to have a label, but we need to do some refactoring for that.
        self.label = str(vlan_no)
@no_dry_run
def apply_networking(net_map):
    def set_access_vlan(port, vlan_id):
        """Set a port to access a given vlan.

        This function expects to be called while in the config prompt, and
        leaves you there when done.
        """
        console.sendline('int ethernet 1/%s' % port)
        # set it first to switch mode then to access mode:
        console.expect(r'[\r\n]+.+# ')
        console.sendline('sw')
        console.expect(r'[\r\n]+.+# ')
        console.sendline('sw mode access')
        console.expect(r'[\r\n]+.+# ')

        if vlan_id is None:
            # turn the port off
            console.sendline('no sw')
        else:
            # set the vlan:
            console.sendline('sw access vlan %s' % vlan_id)
        console.expect(r'[\r\n]+.+# ')
        # back out to config_prompt
        console.sendline('exit')
        console.expect(r'[\r\n]+.+# ')

    nexus_ip = cfg.get('switch nexus', 'ip')
    nexus_user = cfg.get('switch nexus', 'user')
    nexus_pass = cfg.get('switch nexus', 'pass')
    try:
       console = pexpect.spawn('telnet ' + nexus_ip)
       console.expect('login: ')
       console.sendline(nexus_user)
       console.expect('password: ')
       console.sendline(nexus_pass)
       console.expect(r'[\r\n]+.+# ')
       console.sendline('config terminal')
       console.expect(r'[\r\n]+.+# ')

       for port_id in net_map:
            set_access_vlan(port_id, net_map[port_id])

    except IOError as e:
        print "Connection error while connecting to the Nexus switch({0}): {1}".format(e.errno, e.strerror)
    except:
        print "Unexpected error while connecting to the Nexus switch:", sys.exc_info()[0]
        raise
        sys.Exit(1)

def get_new_network_id(db):
    vlan = db.query(Nexus_Vlan).filter_by(available = True).first()
    if not vlan:
        return None
    vlan.available = False
    returnee = str(vlan.vlan_no)
    return returnee

def free_network_id(db, net_id):
    vlan = db.query(Nexus_Vlan).filter_by(vlan_no = net_id).first()
    if not vlan:
        pass
    vlan.available = True

def get_vlan_list():
    vlan_str = cfg.get('switch nexus', 'vlans')
    returnee = []
    for r in vlan_str.split(","):
        r = r.strip().split("-")
        if len(r) == 1:
            returnee.append(int(r[0]))
        else:
            returnee += range(int(r[0]), int(r[1])+1)
    return returnee

def init_db(create=False):
    if not create:
        return
    vlan_list = get_vlan_list()
    db = Session()
    for vlan in vlan_list:
        db.add(Nexus_Vlan(vlan))
    db.commit()

<commit_msg>new line added, in the comments DELLis changed to Nexus<commit_after># Copyright 2013-2014 Massachusetts Open Cloud Contributors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the
# License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an "AS
# IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
# express or implied. See the License for the specific language
# governing permissions and limitations under the License.
"""A switch driver for the Nexus 5500
See the documentation for the haas.drivers package for a description of this
module's interface.
Currently the driver uses telnet to connect to the switch's console; in
the long term we want to be using SNMP.
"""
import os
import pexpect
import re
from haas.config import cfg
from haas.dev_support import no_dry_run
from haas.model import Model, Session
from sqlalchemy import *
class Nexus_Vlan(Model):
    """A VLAN for the Cisco Nexus switch
       This is used to track which vlan numbers are available; when a Network is
       created, it must allocate a Vlan, to ensure that:
       1. The VLAN number it is using is unique, and
       2. The VLAN number is actually allocated to the HaaS; on some deployments we
       may have specific vlan numbers that we are allowed to use.
       This code currently only works if there is only one Nexus switch.
       This is because the Nexus switch does not support switch stacking.
    """
    vlan_no = Column(Integer, nullable=False, unique=True)
    available = Column(Boolean, nullable=False)
    def __init__(self, vlan_no):
        self.vlan_no = vlan_no
        self.available = True
        # XXX: This is pretty gross; it arguably doesn't even make sense for
        # Vlan to have a label, but we need to do some refactoring for that.
        self.label = str(vlan_no)
@no_dry_run
def apply_networking(net_map):
    def set_access_vlan(port, vlan_id):
        """Set a port to access a given vlan.

        This function expects to be called while in the config prompt, and
        leaves you there when done.
        """
        console.sendline('int ethernet 1/%s' % port)
        # set it first to switch mode then to access mode:
        console.expect(r'[\r\n]+.+# ')
        console.sendline('sw')
        console.expect(r'[\r\n]+.+# ')
        console.sendline('sw mode access')
        console.expect(r'[\r\n]+.+# ')

        if vlan_id is None:
            # turn the port off
            console.sendline('no sw')
        else:
            # set the vlan:
            console.sendline('sw access vlan %s' % vlan_id)
        console.expect(r'[\r\n]+.+# ')
        # back out to config_prompt
        console.sendline('exit')
        console.expect(r'[\r\n]+.+# ')

    nexus_ip = cfg.get('switch nexus', 'ip')
    nexus_user = cfg.get('switch nexus', 'user')
    nexus_pass = cfg.get('switch nexus', 'pass')
    try:
       console = pexpect.spawn('telnet ' + nexus_ip)
       console.expect('login: ')
       console.sendline(nexus_user)
       console.expect('password: ')
       console.sendline(nexus_pass)
       console.expect(r'[\r\n]+.+# ')
       console.sendline('config terminal')
       console.expect(r'[\r\n]+.+# ')

       for port_id in net_map:
            set_access_vlan(port_id, net_map[port_id])

    except IOError as e:
        print "Connection error while connecting to the Nexus switch({0}): {1}".format(e.errno, e.strerror)
    except:
        print "Unexpected error while connecting to the Nexus switch:", sys.exc_info()[0]
        raise
        sys.Exit(1)

def get_new_network_id(db):
    vlan = db.query(Nexus_Vlan).filter_by(available = True).first()
    if not vlan:
        return None
    vlan.available = False
    returnee = str(vlan.vlan_no)
    return returnee

def free_network_id(db, net_id):
    vlan = db.query(Nexus_Vlan).filter_by(vlan_no = net_id).first()
    if not vlan:
        pass
    vlan.available = True

def get_vlan_list():
    vlan_str = cfg.get('switch nexus', 'vlans')
    returnee = []
    for r in vlan_str.split(","):
        r = r.strip().split("-")
        if len(r) == 1:
            returnee.append(int(r[0]))
        else:
            returnee += range(int(r[0]), int(r[1])+1)
    return returnee

def init_db(create=False):
    if not create:
        return
    vlan_list = get_vlan_list()
    db = Session()
    for vlan in vlan_list:
        db.add(Nexus_Vlan(vlan))
    db.commit()

<|endoftext|>


<commit_before>from fabric.api import local

vim_bundles = [
    {
        'git': 'git://github.com/mileszs/ack.vim.git',
        'path': '~/.vim/bundle/ack.vim'
    },
    {
        'git': 'git://github.com/tpope/vim-fugitive.git',
        'path': '~/.vim/bundle/vim-fugitive'
    },
    {
        'git': 'git://github.com/sjl/gundo.vim.git',
        'path': '~/.vim/bundle/gundo.vim'
    },
    {
        'git': 'git://github.com/ervandew/supertab.git',
        'path': '~/.vim/bundle/supertab'
    },
    {
        'git': 'git://github.com/Raimondi/delimitMate.git',
        'path': '~/.vim/bundle/delimitMate'
    },
    {
        'git': 'git://github.com/docunext/closetag.vim.git',
        'path': '~/.vim/bundle/closetag.vim'
    },
    {
        'git': 'git://github.com/majutsushi/tagbar.git',
        'path': '~/.vim/bundle/tagbar'
    },
    {
	'git': 'git://github.com/scrooloose/syntastic.git',
	'path': '~/.vim/bundle/syntastic'
    }
]

def apt_get():
    local('sudo apt-get update')
    local('sudo apt-get upgrade')
    local('sudo apt-get install zsh vim wget curl rxvt-unicode suckless-tools xautolock feh tmux mutt mpd ncmpcpp vlc p7zip unrar unzip keepassx xdotool xclip rtorrent diffpdf')

# Not using Haskell at the moment
#def cabal():
#    local('cabal update')
#    local('sudo cabal install yeganesh --global')
#    local('sudo cabal install hoogle hsenv ghc-mod hdevtools hlint pointfree')

def oh_my_zsh():
    local('curl -L http://install.ohmyz.sh | sh')
    local('cp ~/.zshrc.pre-oh-my-zsh ~/.zshrc')
    local('chsh -s $(which shell)')

def install_vim():
    local('mkdir -p ~/.vim/autoload ~/.vim/bundle')
    local('curl -LSso ~/.vim/autoload/pathogen.vim https://tpo.pe/pathogen.vim')

    for bundle in vim_bundles:
        local('git clone ' + bundle['git'] + ' ' + bundle['path'])

    local('cd ~')

def update_vim():
    for bundle in vim_bundles:
        local('cd ' + bundle['path'] + ' && git pull')

    local('cd ~')
<commit_msg>Add xfce4<commit_after>from fabric.api import local

vim_bundles = [
    {
        'git': 'git://github.com/mileszs/ack.vim.git',
        'path': '~/.vim/bundle/ack.vim'
    },
    {
        'git': 'git://github.com/tpope/vim-fugitive.git',
        'path': '~/.vim/bundle/vim-fugitive'
    },
    {
        'git': 'git://github.com/sjl/gundo.vim.git',
        'path': '~/.vim/bundle/gundo.vim'
    },
    {
        'git': 'git://github.com/ervandew/supertab.git',
        'path': '~/.vim/bundle/supertab'
    },
    {
        'git': 'git://github.com/Raimondi/delimitMate.git',
        'path': '~/.vim/bundle/delimitMate'
    },
    {
        'git': 'git://github.com/docunext/closetag.vim.git',
        'path': '~/.vim/bundle/closetag.vim'
    },
    {
        'git': 'git://github.com/majutsushi/tagbar.git',
        'path': '~/.vim/bundle/tagbar'
    },
    {
	'git': 'git://github.com/scrooloose/syntastic.git',
	'path': '~/.vim/bundle/syntastic'
    }
]

def apt_get():
    local('sudo apt-get update')
    local('sudo apt-get upgrade')
    local('sudo apt-get install zsh vim wget curl rxvt-unicode suckless-tools xautolock feh tmux mutt mpd ncmpcpp vlc p7zip unrar unzip keepassx xdotool xclip rtorrent diffpdf xfce4')

# Not using Haskell at the moment
#def cabal():
#    local('cabal update')
#    local('sudo cabal install yeganesh --global')
#    local('sudo cabal install hoogle hsenv ghc-mod hdevtools hlint pointfree')

def oh_my_zsh():
    local('curl -L http://install.ohmyz.sh | sh')
    local('cp ~/.zshrc.pre-oh-my-zsh ~/.zshrc')
    local('chsh -s $(which shell)')

def install_vim():
    local('mkdir -p ~/.vim/autoload ~/.vim/bundle')
    local('curl -LSso ~/.vim/autoload/pathogen.vim https://tpo.pe/pathogen.vim')

    for bundle in vim_bundles:
        local('git clone ' + bundle['git'] + ' ' + bundle['path'])

    local('cd ~')

def update_vim():
    for bundle in vim_bundles:
        local('cd ' + bundle['path'] + ' && git pull')

    local('cd ~')
<|endoftext|>