yandex-tank/Tank/ConsoleWorker.py

""" Provides classes to run TankCore from console environment """
import datetime
import fnmatch
import logging
import os
import sys
import time
import traceback
import signal
from optparse import OptionParser

from Tank.Plugins.ConsoleOnline import RealConsoleMarkup
from tankcore import TankCore


class SingleLevelFilter(logging.Filter):
    """Exclude or approve one msg type at a time.    """

    def __init__(self, passlevel, reject):
        logging.Filter.__init__(self)
        self.passlevel = passlevel
        self.reject = reject

    def filter(self, record):
        if self.reject:
            return record.levelno != self.passlevel
        else:
            return record.levelno == self.passlevel


def signal_handler(sig, frame):
    """ required for non-tty python runs to interrupt """
    raise KeyboardInterrupt()


signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)


class ConsoleTank:
    """    Worker class that runs tank core accepting cmdline params    """

    IGNORE_LOCKS = "ignore_locks"

    def __init__(self, options, ammofile):
        self.core = TankCore()

        self.options = options
        self.ammofile = ammofile

        self.baseconfigs_location = '/etc/yandex-tank'

        self.log_filename = self.options.log
        self.core.add_artifact_file(self.log_filename)
        self.log = logging.getLogger(__name__)

        self.signal_count = 0
        self.scheduled_start = None


    def set_baseconfigs_dir(self, directory):
        """        Set directory where to read configs set        """
        self.baseconfigs_location = directory


    def init_logging(self):
        """        Set up logging, as it is very important for console tool        """
        logger = logging.getLogger()
        logger.setLevel(logging.DEBUG)

        # create file handler which logs even debug messages
        if self.log_filename:
            file_handler = logging.FileHandler(self.log_filename)
            file_handler.setLevel(logging.DEBUG)
            file_handler.setFormatter(logging.Formatter("%(asctime)s [%(levelname)s] %(name)s %(message)s"))
            logger.addHandler(file_handler)

        # create console handler with a higher log level
        console_handler = logging.StreamHandler(sys.stdout)
        stderr_hdl = logging.StreamHandler(sys.stderr)

        fmt_verbose = logging.Formatter("%(asctime)s [%(levelname)s] %(name)s %(message)s")
        fmt_regular = logging.Formatter("%(asctime)s %(levelname)s: %(message)s", "%H:%M:%S")

        if self.options.verbose:
            console_handler.setLevel(logging.DEBUG)
            console_handler.setFormatter(fmt_verbose)
            stderr_hdl.setFormatter(fmt_verbose)
        elif self.options.quiet:
            console_handler.setLevel(logging.WARNING)
            console_handler.setFormatter(fmt_regular)
            stderr_hdl.setFormatter(fmt_regular)
        else:
            console_handler.setLevel(logging.INFO)
            console_handler.setFormatter(fmt_regular)
            stderr_hdl.setFormatter(fmt_regular)

        f_err = SingleLevelFilter(logging.ERROR, True)
        f_warn = SingleLevelFilter(logging.WARNING, True)
        f_crit = SingleLevelFilter(logging.CRITICAL, True)
        console_handler.addFilter(f_err)
        console_handler.addFilter(f_warn)
        console_handler.addFilter(f_crit)
        logger.addHandler(console_handler)

        f_info = SingleLevelFilter(logging.INFO, True)
        f_debug = SingleLevelFilter(logging.DEBUG, True)
        stderr_hdl.addFilter(f_info)
        stderr_hdl.addFilter(f_debug)
        logger.addHandler(stderr_hdl)


    def __override_config_from_cmdline(self):
        """ override config options from command line"""
        if self.options.option:
            self.core.apply_shorthand_options(self.options.option)


    def get_default_configs(self):
        """ returns default configs list, from /etc and home dir """
        configs = []
        try:
            conf_files = os.listdir(self.baseconfigs_location)
            conf_files.sort()
            for filename in conf_files:
                if fnmatch.fnmatch(filename, '*.ini'):
                    configs += [os.path.realpath(self.baseconfigs_location + os.sep + filename)]
        except OSError:
            self.log.warn(self.baseconfigs_location + ' is not acessible to get configs list')

        configs += [os.path.expanduser('~/.yandex-tank')]
        return configs


    def configure(self):
        """         Make all console-specific preparations before running Tank        """
        if self.options.ignore_lock:
            self.log.warn("Lock files ignored. This is highly unrecommended practice!")

        if self.options.lock_dir:
            self.core.set_option(self.core.SECTION, "lock_dir", self.options.lock_dir)

        while True:
            try:
                self.core.get_lock(self.options.ignore_lock)
                break
            except Exception, exc:
                if self.options.lock_fail:
                    raise RuntimeError("Lock file present, cannot continue")
                self.log.debug("Failed to get lock: %s", traceback.format_exc(exc))
                self.log.info("Waiting 5s for retry...")
                time.sleep(5)

        try:
            configs = []

            if not self.options.no_rc:
                configs = self.get_default_configs()

            if not self.options.config:
                if os.path.exists(os.path.realpath('load.ini')):
                    self.log.info("No config passed via cmdline, using ./load.ini")
                    configs += [os.path.realpath('load.ini')]
                    self.core.add_artifact_file(os.path.realpath('load.ini'), True)
                elif os.path.exists(os.path.realpath('load.conf')):
                    # just for old 'lunapark' compatibility
                    self.log.warn("Using 'load.conf' is unrecommended, please use 'load.ini' instead")
                    conf_file = os.path.realpath('load.conf')
                    configs += [conf_file]
                    self.core.add_artifact_file(conf_file, True)
            else:
                for config_file in self.options.config:
                    configs.append(config_file)

            self.core.load_configs(configs)

            if self.ammofile:
                self.log.debug("Ammofile: %s", self.ammofile)
                self.core.set_option("phantom", 'ammofile', self.ammofile[0])

            self.__override_config_from_cmdline()

            self.core.load_plugins()

            if self.options.scheduled_start:
                try:
                    self.scheduled_start = datetime.datetime.strptime(self.options.scheduled_start, '%Y-%m-%d %H:%M:%S')
                except ValueError:
                    self.scheduled_start = datetime.datetime.strptime(
                        datetime.datetime.now().strftime('%Y-%m-%d ') + self.options.scheduled_start,
                        '%Y-%m-%d %H:%M:%S')

            if self.options.ignore_lock:
                self.core.set_option(self.core.SECTION, self.IGNORE_LOCKS, "1")

        except Exception, ex:
            self.log.info("Exception: %s", traceback.format_exc(ex))
            sys.stderr.write(RealConsoleMarkup.RED)
            self.log.error("%s", ex)
            sys.stderr.write(RealConsoleMarkup.RESET)
            sys.stderr.write(RealConsoleMarkup.TOTAL_RESET)
            self.core.release_lock()
            raise ex

    def __graceful_shutdown(self):
        """ call shutdown routines """
        retcode = 1
        self.log.info("Trying to shutdown gracefully...")
        retcode = self.core.plugins_end_test(retcode)
        retcode = self.core.plugins_post_process(retcode)
        self.log.info("Done graceful shutdown")
        return retcode


    def perform_test(self):
        """
        Run the test sequence via Tank Core
        """
        self.log.info("Performing test")
        retcode = 1
        try:
            self.core.plugins_configure()
            self.core.plugins_prepare_test()
            if self.scheduled_start:
                self.log.info("Waiting scheduled time: %s...", self.scheduled_start)
                while datetime.datetime.now() < self.scheduled_start:
                    self.log.debug("Not yet: %s < %s", datetime.datetime.now(), self.scheduled_start)
                    time.sleep(1)
                self.log.info("Time has come: %s", datetime.datetime.now())

            if self.options.manual_start:
                raw_input("Press Enter key to start test:")

            self.core.plugins_start_test()
            retcode = self.core.wait_for_finish()
            retcode = self.core.plugins_end_test(retcode)
            retcode = self.core.plugins_post_process(retcode)

        except KeyboardInterrupt as ex:
            sys.stdout.write(RealConsoleMarkup.YELLOW)
            self.log.info("Do not press Ctrl+C again, the test will be broken otherwise")
            sys.stdout.write(RealConsoleMarkup.RESET)
            sys.stdout.write(RealConsoleMarkup.TOTAL_RESET)
            self.signal_count += 1
            self.log.debug("Caught KeyboardInterrupt: %s", traceback.format_exc(ex))
            try:
                retcode = self.__graceful_shutdown()
            except KeyboardInterrupt as ex:
                self.log.debug("Caught KeyboardInterrupt again: %s", traceback.format_exc(ex))
                self.log.info("User insists on exiting, aborting graceful shutdown...")
                retcode = 1

        except Exception as ex:
            self.log.info("Exception: %s", traceback.format_exc(ex))
            sys.stderr.write(RealConsoleMarkup.RED)
            self.log.error("%s", ex)
            sys.stderr.write(RealConsoleMarkup.RESET)
            sys.stderr.write(RealConsoleMarkup.TOTAL_RESET)
            retcode = self.__graceful_shutdown()
            self.core.release_lock()

        self.log.info("Done performing test with code %s", retcode)
        return retcode


class DevNullOpts:
    def __init__(self):
        pass

    log = "/dev/null"


class CompletionHelperOptionParser(OptionParser):
    def __init__(self):
        OptionParser.__init__(self, add_help_option=False)
        self.add_option('--bash-switches-list', action='store_true', dest="list_switches", help="Options list")
        self.add_option('--bash-options-prev', action='store', dest="list_options_prev", help="Options list")
        self.add_option('--bash-options-cur', action='store', dest="list_options_cur", help="Options list")

    def error(self, msg):
        pass

    def exit(self, status=0, msg=None):
        pass

    def handle_request(self, parser):
        options = self.parse_args()[0]
        if options.list_switches:
            opts = []
            for option in parser.option_list:
                if not "--bash" in option.get_opt_string():
                    opts.append(option.get_opt_string())
            print ' '.join(opts)
            exit(0)

        if options.list_options_cur or options.list_options_prev:
            cmdtank = ConsoleTank(DevNullOpts(), None)
            cmdtank.core.load_configs(cmdtank.get_default_configs())
            cmdtank.core.load_plugins()

            opts = []
            for option in cmdtank.core.get_available_options():
                opts.append(cmdtank.core.SECTION + '.' + option + '=')

            plugin_keys = cmdtank.core.config.get_options(cmdtank.core.SECTION, cmdtank.core.PLUGIN_PREFIX)
            for (plugin_name, plugin_path) in plugin_keys:
                opts.append(cmdtank.core.SECTION + '.' + cmdtank.core.PLUGIN_PREFIX + plugin_name + '=')

            for plugin in cmdtank.core.plugins.values():
                for option in plugin.get_available_options():
                    opts.append(plugin.SECTION + '.' + option + '=')
            print ' '.join(sorted(opts))
            exit(0)