~veebers/juju-ci-tools/model_migration_check_all_units_of_charm

#!/usr/bin/env python

"""This module tests the deployment with constraints."""

from __future__ import print_function

import argparse

import logging

import os

import sys

import re

from deploy_stack import (

    BootstrapManager,

    )

from jujucharm import (

    Charm,

    local_charm_path,

    )

from utility import (

    add_basic_testing_arguments,

    configure_logging,

    JujuAssertionError,

    temp_dir,

    )

__metaclass__ = type

log = logging.getLogger("assess_constraints")

VIRT_TYPES = ['lxd']

INSTANCE_TYPES = {

    'azure': [],

    'ec2': ['t2.micro'],

    'gce': [],

    'joyent': [],

    'openstack': [],

    }

# This assumes instances are unique accross providers.

def get_instance_spec(instance_type):

    """Get the specifications of a given instance type."""

    return {

        't2.micro': {'mem': '1G', 'cpu-power': '10', 'cores': '1'},

        }[instance_type]

def mem_to_int(size):

    """Convert an argument size into a number of megabytes."""

    if not re.match(re.compile('^[0123456789]+[MGTP]?$'), size):

        raise JujuAssertionError('Not a size format:', size)

    if size[-1] in 'MGTP':

        val = int(size[0:-1])

        unit = size[-1]

        return val * (1024 ** 'MGTP'.find(unit))

    else:

        return int(size)

class Constraints:

    """Class that represents a set of contraints."""

    @staticmethod

    def _list_to_str(constraints_list):

        parts = ['{}={}'.format(name, value) for (name, value) in

                 constraints_list if value is not None]

        return ' '.join(parts)

    def _get_constraint_pairs(self):

        """Get a list of (constraint-name, constraint-value) pairs."""

        return [('mem', self.mem), ('cores', self.cores),

                ('virt-type', self.virt_type),

                ('instance-type', self.instance_type),

                ('root-disk', self.root_disk), ('cpu-power', self.cpu_power),

                ('arch', self.arch),

                ]

    def __init__(self, mem=None, cores=None, virt_type=None,

                 instance_type=None, root_disk=None, cpu_power=None,

                 arch=None):

        """Create a new constraints instance from individual constraints."""

        self.mem = mem

        self.cores = cores

        self.virt_type = virt_type

        self.instance_type = instance_type

        self.root_disk = root_disk

        self.cpu_power = cpu_power

        self.arch = arch

    def __repr__(self):

        """Get a detailed string reperentation of the object."""

        pairs = self._get_constraint_pairs()

        parts = ['{}={!r}'.format(name.replace('-', '_'), value)

                 for (name, value) in pairs if value is not None]

        return 'Constraints({})'.format(', '.join(parts))

    def __str__(self):

        """Convert the instance constraint values into an argument string."""

        return Constraints._list_to_str(self._get_constraint_pairs())

    def __eq__(self, other):

        return (self.mem == other.mem and self.cores == other.cores and

                self.virt_type == other.virt_type and

                self.instance_type == other.instance_type and

                self.root_disk == other.root_disk and

                self.cpu_power == other.cpu_power and

                self.arch == other.arch

                )

    @staticmethod

    def _meets_string(constraint, actual):

        if constraint is None:

            return True

        if actual is None:

            return False

        return constraint == actual

    @staticmethod

    def _meets_min_int(constraint, actual):

        if constraint is None:

            return True

        if actual is None:

            return False

        return int(constraint) <= int(actual)

    @staticmethod

    def _meets_min_mem(constraint, actual):

        if constraint is None:

            return True

        if actual is None:

            return False

        return mem_to_int(constraint) <= mem_to_int(actual)

    def meets_root_disk(self, actual_root_disk):

        """Check to see if a given value meets the root_disk constraint."""

        return self._meets_min_mem(self.root_disk, actual_root_disk)

    def meets_cores(self, actual_cores):

        """Check to see if a given value meets the cores constraint."""

        return self._meets_min_int(self.cores, actual_cores)

    def meets_cpu_power(self, actual_cpu_power):

        """Check to see if a given value meets the cpu_power constraint."""

        return self._meets_min_int(self.cpu_power, actual_cpu_power)

    def meets_arch(self, actual_arch):

        """Check to see if a given value meets the arch constraint."""

        return self._meets_string(self.arch, actual_arch)

    def meets_instance_type(self, actual_data):

        """Check to see if a given value meets the instance_type constraint.

        Currently there is no direct way to check for it, so we 'fingerprint'

        each instance_type in a dictionary."""

        if self.instance_type is None:

            return True

        instance_data = get_instance_spec(self.instance_type)

        for (key, value) in instance_data.iteritems():

            # Temperary fix until cpu-cores -> cores switch is finished.

            if key is 'cores' and 'cpu-cores' in actual_data:

                key = 'cpu-cores'

            if key not in actual_data:

                raise JujuAssertionError('Missing data:', key)

            elif key in ['mem', 'root-disk']:

                if mem_to_int(value) != mem_to_int(actual_data[key]):

                    return False

            elif value != actual_data[key]:

                return False

        else:

            return True

    def meets_all(self, actual_data):

        return (self.meets_root_disk(actual_data.get('root-disk')) and

                self.meets_cores(actual_data.get('cores')) and

                self.meets_cpu_power(actual_data.get('cpu-power')) and

                self.meets_arch(actual_data.get('arch')) and

                self.meets_instance_type(actual_data))

def deploy_constraint(client, constraints, charm, series, charm_repo):

    """Test deploying charm with constraints."""

    client.deploy(charm, series=series, repository=charm_repo,

                  constraints=str(constraints))

    client.wait_for_workloads()

def deploy_charm_constraint(client, constraints, charm_name, charm_series,

                            charm_dir):

    """Create a charm with constraints and test deploying it."""

    constraints_charm = Charm(charm_name,

                              'Test charm for constraints',

                              series=[charm_series])

    charm_root = constraints_charm.to_repo_dir(charm_dir)

    platform = 'ubuntu'

    charm = local_charm_path(charm=charm_name,

                             juju_ver=client.version,

                             series=charm_series,

                             repository=os.path.dirname(charm_root),

                             platform=platform)

    deploy_constraint(client, constraints, charm,

                      charm_series, charm_dir)

def machine_hardware(client, machine):

    """Get hardware data about the given machine."""

    machine_data = client.show_machine(machine)

    hardware = machine_data['machines'][machine]['hardware']

    data = {}

    for kvp in hardware.split(' '):

        (key, value) = kvp.split('=')

        data[key] = value

    return data

def application_machines(client, application):

    """Get all the machines used to host the given application."""

    status = client.get_status()

    app_data = status.get_applications()[application]

    machines = [unit_data['machine'] for unit_data in

                app_data['units'].values()]

    return machines

def application_hardware(client, application):

    """Get hardware data about a machine for an application."""

    machines = application_machines(client, application)

    return machine_hardware(client, machines[0])

def prepare_constraint_test(client, constraints, charm_name,

                            charm_series='xenial'):

    """Deploy a charm with constraints and data to see if it meets them."""

    with temp_dir() as charm_dir:

        deploy_charm_constraint(client, constraints, charm_name,

                                charm_series, charm_dir)

        client.wait_for_started()

        return application_hardware(client, charm_name)

def assess_virt_type(client, virt_type):

    """Assess the virt-type option for constraints"""

    if virt_type not in VIRT_TYPES:

        raise JujuAssertionError(virt_type)

    constraints = Constraints(virt_type=virt_type)

    charm_name = 'virt-type-{}'.format(virt_type)

    charm_series = 'xenial'

    with temp_dir() as charm_dir:

        deploy_charm_constraint(client, constraints, charm_name,

                                charm_series, charm_dir)

def assess_virt_type_constraints(client, test_kvm=False):

    """Assess deployment with virt-type constraints."""

    if test_kvm:

        VIRT_TYPES.append("kvm")

    for virt_type in VIRT_TYPES:

        assess_virt_type(client, virt_type)

    try:

        assess_virt_type(client, 'aws')

    except JujuAssertionError:

        log.info("Correctly rejected virt-type aws")

    else:

        raise JujuAssertionError("FAIL: Client deployed with virt-type aws")

    if test_kvm:

        VIRT_TYPES.remove("kvm")

def get_failure_exception(client, constraints):

    """Create a JujuAssertionError with a detailed error message."""

    message = 'Test Failed: on {} with constraints "{}"'.format(

        client.env.provider, str(constraints))

    return JujuAssertionError(message)

def assess_constraints_deploy(client, constraints, charm_name):

    """Check a single set of constraints on deploy.

    :param client: Client to deploy the charm to.

    :param constraints: Constraints used and checked against.

    :param charm_name: Name of the charm to try deploying.

    :raises JujuAssertionError if test fails."""

    data = prepare_constraint_test(client, constraints, charm_name)

    if not constraints.meets_all(data):

        raise get_failure_exception(client, constraints)

def assess_instance_type(client, provider, instance_type):

    """Assess the instance-type option for constraints"""

    if instance_type not in INSTANCE_TYPES[provider]:

        raise JujuAssertionError(instance_type)

    constraints = Constraints(instance_type=instance_type)

    charm_name = 'instance-type-{}'.format(instance_type.replace('.', '-'))

    assess_constraints_deploy(client, constraints, charm_name)

def assess_instance_type_constraints(client, provider=None):

    """Assess deployment with instance-type constraints."""

    if provider is None:

        provider = client.env.provider

    if provider not in INSTANCE_TYPES:

        return

    for instance_type in INSTANCE_TYPES[provider]:

        assess_instance_type(client, provider, instance_type)

def assess_root_disk_constraints(client, values):

    """Assess deployment with root-disk constraints."""

    for root_disk in values:

        constraints = Constraints(root_disk=root_disk)

        charm_name = 'root-disk-{}'.format(root_disk.lower())

        assess_constraints_deploy(client, constraints, charm_name)

def assess_cores_constraints(client, values):

    """Assess deployment with cores constraints."""

    for cores in values:

        constraints = Constraints(cores=cores)

        charm_name = 'cores-{}c'.format(cores)

        assess_constraints_deploy(client, constraints, charm_name)

def assess_cpu_power_constraints(client, values):

    """Assess deployment with cpu_power constraints."""

    for cpu_power in values:

        constraints = Constraints(cpu_power=cpu_power)

        charm_name = 'cpu-power-{}cp'.format(cpu_power)

        assess_constraints_deploy(client, constraints, charm_name)

def assess_multiple_constraints(client, base_name, **kwargs):

    """Assess deployment with muliple_constraints.

    Makes sure the combination of constraints gives us new instance type."""

    finger_prints = []

    for (part, (constraint, value)) in enumerate(kwargs.iteritems()):

        data = prepare_constraint_test(

            client, Constraints(**{constraint: value}),

            '{}-part{}'.format(base_name, part))

        finger_prints.append(data)

    final_constraints = Constraints(**kwargs)

    data = prepare_constraint_test(client, final_constraints,

                                   '{}-whole'.format(base_name))

    if not final_constraints.meets_all(data):

        raise get_failure_exception(client, final_constraints)

    if data in finger_prints:

        raise JujuAssertionError(

            'Multiple Constraints did not change the hardware.')

def assess_constraints(client, test_kvm=False):

    """Assess deployment with constraints."""

    provider = client.env.provider

    if 'lxd' == provider:

        assess_virt_type_constraints(client, test_kvm)

    elif 'ec2' == provider:

        assess_instance_type_constraints(client, provider)

        assess_root_disk_constraints(client, ['16G'])

        assess_cores_constraints(client, ['2'])

        assess_cpu_power_constraints(client, ['30'])

        assess_multiple_constraints(client, 'root-disk-and-cpu-power',

                                    root_disk='15G', cpu_power='40')

def parse_args(argv):

    """Parse all arguments."""

    parser = argparse.ArgumentParser(description="Test constraints")

    add_basic_testing_arguments(parser)

    return parser.parse_args(argv)

def main(argv=None):

    args = parse_args(argv)

    configure_logging(args.verbose)

    bs_manager = BootstrapManager.from_args(args)

    test_kvm = '--with-virttype-kvm' in args

    with bs_manager.booted_context(args.upload_tools):

        assess_constraints(bs_manager.client, test_kvm)

    return 0

if __name__ == '__main__':

    sys.exit(main())