github.com/niedbalski/juju@v0.0.0-20190215020005-8ff100488e47/acceptancetests/utility.py

from contextlib import contextmanager
from datetime import (
    datetime,
    timedelta,
)
import errno
import json
import logging
import os
import re
import subprocess
import socket
import sys
from time import (
    sleep,
    time,
)
from jujupy.utility import (
    ensure_deleted,
    ensure_dir,
    get_timeout_path,
    get_unit_public_ip,
    is_ipv6_address,
    print_now,
    qualified_model_name,
    quote,
    scoped_environ,
    skip_on_missing_file,
    temp_dir,
    temp_yaml_file,
    until_timeout
)

# Imported for other call sites to use.
__all__ = [
    'ensure_deleted',
    'ensure_dir',
    'get_timeout_path',
    'get_unit_public_ip',
    'qualified_model_name',
    'quote',
    'scoped_environ',
    'skip_on_missing_file',
    'temp_dir',
    'temp_yaml_file',
]


# Equivalent of socket.EAI_NODATA when using windows sockets
# <https://msdn.microsoft.com/ms740668#WSANO_DATA>
WSANO_DATA = 11004

TEST_MODEL = 'test-tmp-env'

log = logging.getLogger("utility")


class PortTimeoutError(Exception):
    pass


class LoggedException(BaseException):
    """Raised in place of an exception that has already been logged.

    This is a wrapper to avoid double-printing real Exceptions while still
    unwinding the stack appropriately.
    """

    def __init__(self, exception):
        self.exception = exception


class JujuAssertionError(AssertionError):
    """Exception for juju assertion failures."""


def _clean_dir(maybe_dir):
    """Pseudo-type that validates an argument to be a clean directory path.

    For safety, this function will not attempt to remove existing directory
    contents but will just report a warning.
    """
    try:
        contents = os.listdir(maybe_dir)
    except OSError as e:
        if e.errno == errno.ENOENT:
            # we don't raise this error due to tests abusing /tmp/logs
            logging.warning('Not a directory {}'.format(maybe_dir))
        if e.errno == errno.EEXIST:
            logging.warnings('Directory {} already exists'.format(maybe_dir))
    else:
        if contents and contents != ["empty"]:
            logging.warning(
                'Directory {!r} has existing contents.'.format(maybe_dir))
    return maybe_dir


def as_literal_address(address):
    """Returns address in form suitable for embedding in URL or similar.

    In practice, this just puts square brackets round IPv6 addresses which
    avoids conflict with port seperators and other uses of colons.
    """
    if is_ipv6_address(address):
        return address.join("[]")
    return address


def wait_for_port(host, port, closed=False, timeout=30):
    family = socket.AF_INET6 if is_ipv6_address(host) else socket.AF_INET
    for remaining in until_timeout(timeout):
        try:
            addrinfo = socket.getaddrinfo(host, port, family,
                                          socket.SOCK_STREAM)
        except socket.error as e:
            if e.errno not in (socket.EAI_NODATA, WSANO_DATA):
                raise
            if closed:
                return
            else:
                continue
        sockaddr = addrinfo[0][4]
        # Treat Azure messed-up address lookup as a closed port.
        if sockaddr[0] == '0.0.0.0':
            if closed:
                return
            else:
                continue
        conn = socket.socket(*addrinfo[0][:3])
        conn.settimeout(max(remaining or 0, 5))
        try:
            conn.connect(sockaddr)
        except socket.timeout:
            if closed:
                return
        except socket.error as e:
            if e.errno not in (errno.ECONNREFUSED, errno.ENETUNREACH,
                               errno.ETIMEDOUT, errno.EHOSTUNREACH):
                raise
            if closed:
                return
        except socket.gaierror as e:
            print_now(str(e))
        except Exception as e:
            print_now('Unexpected {!r}: {}'.format((type(e), e)))
            raise
        else:
            conn.close()
            if not closed:
                return
            sleep(1)
    raise PortTimeoutError('Timed out waiting for port.')


def get_revision_build(build_info):
    for action in build_info['actions']:
        if 'parameters' in action:
            for parameter in action['parameters']:
                if parameter['name'] == 'revision_build':
                    return parameter['value']


def get_winrm_certs():
    """"Returns locations of key and cert files for winrm in cloud-city."""
    home = os.environ['HOME']
    return (
        os.path.join(home, 'cloud-city/winrm_client_cert.key'),
        os.path.join(home, 'cloud-city/winrm_client_cert.pem'),
    )


def s3_cmd(params, drop_output=False):
    s3cfg_path = os.path.join(
        os.environ['HOME'], 'cloud-city/juju-qa.s3cfg')
    command = ['s3cmd', '-c', s3cfg_path, '--no-progress'] + params
    if drop_output:
        return subprocess.check_call(
            command, stdout=open('/dev/null', 'w'))
    else:
        return subprocess.check_output(command)


def _get_test_name_from_filename():
    try:
        calling_file = sys._getframe(2).f_back.f_globals['__file__']
        return os.path.splitext(os.path.basename(calling_file))[0]
    except:
        return 'unknown_test'


def generate_default_clean_dir(temp_env_name):
    """Creates a new unique directory for logging and returns name"""
    logging.debug('Environment {}'.format(temp_env_name))
    test_name = temp_env_name.split('-')[0]
    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
    log_dir = os.path.join('/tmp', test_name, 'logs', timestamp)

    try:
        os.makedirs(log_dir)
        logging.info('Created logging directory {}'.format(log_dir))
    except OSError as e:
        if e.errno == errno.EEXIST:
            logging.warn('"Directory {} already exists'.format(log_dir))
        else:
            raise('Failed to create logging directory: {} ' +
                  log_dir +
                  '. Please specify empty folder or try again')
    return log_dir


def _generate_default_temp_env_name():
    """Creates a new unique name for environment and returns the name"""
    # we need to sanitize the name
    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
    test_name = re.sub('[^a-zA-Z]', '', _get_test_name_from_filename())
    return '{}-{}-temp-env'.format(test_name, timestamp)


def _to_deadline(timeout):
    return datetime.utcnow() + timedelta(seconds=int(timeout))


def add_arg_juju_bin(parser):
    parser.add_argument('juju_bin', nargs='?',
                        help='Full path to the Juju binary. By default, this'
                        ' will use $PATH/juju',
                        default=None)


def add_basic_testing_arguments(
        parser, using_jes=False, deadline=True, env=True, existing=True):
    """Returns the parser loaded with basic testing arguments.

    The basic testing arguments, used in conjuction with boot_context ensures
    a test can be run in any supported substrate in parallel.

    This helper adds 4 positional arguments that defines the minimum needed
    to run a test script.

    These arguments (env, juju_bin, logs, temp_env_name) allow you to specify
    specifics for which env, juju binary, which folder for logging and an
    environment name for your test respectively.

    There are many optional args that either update the env's config or
    manipulate the juju command line options to test in controlled situations
    or in uncommon substrates: --debug, --verbose, --agent-url, --agent-stream,
    --series, --bootstrap-host, --machine, --keep-env. If not using_jes, the
    --upload-tools arg will also be added.

    :param parser: an ArgumentParser.
    :param using_jes: whether args should be tailored for JES testing.
    :param deadline: If true, support the --timeout option and convert to a
        deadline.
    :param existing: If true will supply the 'existing' argument to allow
        running on an existing bootstrapped controller.
    """

    # Optional postional arguments
    if env:
        parser.add_argument(
            'env', nargs='?',
            help='The juju environment to base the temp test environment on.',
            default='lxd')
    add_arg_juju_bin(parser)
    parser.add_argument('logs', nargs='?', type=_clean_dir,
                        help='A directory in which to store logs. By default,'
                        ' this will use the current directory',
                        default=None)
    parser.add_argument('temp_env_name', nargs='?',
                        help='A temporary test environment name. By default, '
                        ' this will generate an enviroment name using the '
                        ' timestamp and testname. '
                        ' test_name_timestamp_temp_env',
                        default=_generate_default_temp_env_name())

    # Optional keyword arguments.
    parser.add_argument('--debug', action='store_true',
                        help='Pass --debug to Juju.')
    parser.add_argument('--verbose', action='store_const',
                        default=logging.INFO, const=logging.DEBUG,
                        help='Verbose test harness output.')
    parser.add_argument('--region', help='Override environment region.')
    parser.add_argument('--to', default=None,
                        help='Place the controller at a location.')
    parser.add_argument('--agent-url', action='store', default=None,
                        help='URL for retrieving agent binaries.')
    parser.add_argument('--agent-stream', action='store', default=None,
                        help='Stream for retrieving agent binaries.')
    parser.add_argument('--series', action='store', default=None,
                        help='Name of the Ubuntu series to use.')
    if not using_jes:
        parser.add_argument('--upload-tools', action='store_true',
                            help='upload local version of tools to bootstrap.')
    parser.add_argument('--bootstrap-host',
                        help='The host to use for bootstrap.')
    parser.add_argument('--machine', help='A machine to add or when used with '
                        'KVM based MaaS, a KVM image to start.',
                        action='append', default=[])
    parser.add_argument('--keep-env', action='store_true',
                        help='Keep the Juju environment after the test'
                        ' completes.')
    parser.add_argument('--logging-config',
                        help="Override logging configuration for a deployment.",
                        default="<root>=INFO;unit=INFO")
    if existing:
        parser.add_argument(
            '--existing',
            action='store',
            default=None,
            const='current',
            nargs='?',
            help='Test using an existing bootstrapped controller. '
                 'If no controller name is provided defaults to using the '
                 'current selected controller.')
    if deadline:
        parser.add_argument('--timeout', dest='deadline', type=_to_deadline,
                            help="The script timeout, in seconds.")
    return parser


# suppress nosetests
add_basic_testing_arguments.__test__ = False


def configure_logging(log_level):
    logging.basicConfig(
        level=log_level, format='%(asctime)s %(levelname)s %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S')


def get_candidates_path(root_dir):
    return os.path.join(root_dir, 'candidate')


# GZ 2015-10-15: Paths returned in filesystem dependent order, may want sort?
def find_candidates(root_dir, find_all=False):
    return (path for path, buildvars in _find_candidates(root_dir, find_all))


def find_latest_branch_candidates(root_dir):
    """Return a list of one candidate per branch.

    :param root_dir: The root directory to find candidates from.
    """
    candidates = []
    for path, buildvars_path in _find_candidates(root_dir, find_all=False,
                                                 artifacts=True):
        with open(buildvars_path) as buildvars_file:
            buildvars = json.load(buildvars_file)
            candidates.append(
                (buildvars['branch'], int(buildvars['revision_build']), path))
    latest = dict(
        (branch, (path, build)) for branch, build, path in sorted(candidates))
    return latest.values()


def _find_candidates(root_dir, find_all=False, artifacts=False):
    candidates_path = get_candidates_path(root_dir)
    a_week_ago = time() - timedelta(days=7).total_seconds()
    for candidate_dir in os.listdir(candidates_path):
        if candidate_dir.endswith('-artifacts') != artifacts:
            continue
        candidate_path = os.path.join(candidates_path, candidate_dir)
        buildvars = os.path.join(candidate_path, 'buildvars.json')
        try:
            stat = os.stat(buildvars)
        except OSError as e:
            if e.errno in (errno.ENOENT, errno.ENOTDIR):
                continue
            raise
        if not find_all and stat.st_mtime < a_week_ago:
            continue
        yield candidate_path, buildvars


def get_deb_arch():
    """Get the debian machine architecture."""
    return subprocess.check_output(['dpkg', '--print-architecture']).strip()


def extract_deb(package_path, directory):
    """Extract a debian package to a specified directory."""
    subprocess.check_call(['dpkg', '-x', package_path, directory])


def run_command(command, dry_run=False, verbose=False):
    """Optionally execute a command and maybe print the output."""
    if verbose:
        print_now('Executing: {}'.format(command))
    if not dry_run:
        output = subprocess.check_output(command)
        if verbose:
            print_now(output)


def log_and_wrap_exception(logger, exc):
    """Record exc details to logger and return wrapped in LoggedException."""
    logger.exception(exc)
    stdout = getattr(exc, 'output', None)
    stderr = getattr(exc, 'stderr', None)
    if stdout or stderr:
        logger.info('Output from exception:\nstdout:\n%s\nstderr:\n%s',
                    stdout, stderr)
    return LoggedException(exc)


@contextmanager
def logged_exception(logger):
    """\
    Record exceptions in managed context to logger and reraise LoggedException.

    Note that BaseException classes like SystemExit, GeneratorExit and
    LoggedException itself are not wrapped, except for KeyboardInterrupt.
    """
    try:
        yield
    except (Exception, KeyboardInterrupt) as e:
        raise log_and_wrap_exception(logger, e)


def assert_dict_is_subset(sub_dict, super_dict):
    """Assert that every item in the sub_dict is in the super_dict.

    :raises JujuAssertionError: when sub_dict items are missing.
    :return: True when when sub_dict is a subset of super_dict
    """
    if not all(item in super_dict.items() for item in sub_dict.items()):
        raise JujuAssertionError(
            'Found: {} \nExpected: {}'.format(super_dict, sub_dict))
    return True


def add_model(client):
    """Adds a model to the current juju environment then destroys it.

    Will raise an exception if the Juju does not deselect the current model.
    :param client: Jujupy ModelClient object
    """
    log.info('Adding model "{}" to current controller'.format(TEST_MODEL))
    new_client = client.add_model(TEST_MODEL)
    new_model = get_current_model(new_client)
    if new_model == TEST_MODEL:
        log.info('Current model and newly added model match')
    else:
        error = ('Juju failed to switch to new model after creation. '
                 'Expected {} got {}'.format(TEST_MODEL, new_model))
        raise JujuAssertionError(error)
    return new_client


def get_current_model(client):
    """Gets the current model from Juju's list-models command.

    :param client: Jujupy ModelClient object
    :return: String name of current model
    """
    raw = list_models(client)
    try:
        return raw['current-model']
    except KeyError:
        log.warning('No model is currently selected.')
        return None


def list_models(client):
    """List models.
    :param client: Jujupy ModelClient object
    :return: Dict of list-models command
    """
    try:
        raw = client.get_juju_output('list-models', '--format', 'json',
                                     include_e=False).decode('utf-8')
    except subprocess.CalledProcessError as e:
        log.error('Failed to list current models due to error: {}'.format(e))
        raise e
    return json.loads(raw)