Example #1
0
 def test_get_task_ratio_dict_total(self):
     ratio = get_task_ratio_dict([MyTaskSet], total=True)
     self.assertEqual(1.0, ratio["MyTaskSet"]["ratio"])
     self.assertEqual(0.75, ratio["MyTaskSet"]["tasks"]["root_task"]["ratio"])
     self.assertEqual(0.25, ratio["MyTaskSet"]["tasks"]["MySubTaskSet"]["ratio"])
     self.assertEqual(0.125, ratio["MyTaskSet"]["tasks"]["MySubTaskSet"]["tasks"]["task1"]["ratio"])
     self.assertEqual(0.125, ratio["MyTaskSet"]["tasks"]["MySubTaskSet"]["tasks"]["task2"]["ratio"])
Example #2
0
 def test_get_task_ratio_dict_total(self):
     ratio = get_task_ratio_dict([MyTaskSet], total=True)
     self.assertEqual(1.0, ratio["MyTaskSet"]["ratio"])
     self.assertEqual(0.75, ratio["MyTaskSet"]["tasks"]["root_task"]["ratio"])
     self.assertEqual(0.25, ratio["MyTaskSet"]["tasks"]["MySubTaskSet"]["ratio"])
     self.assertEqual(0.125, ratio["MyTaskSet"]["tasks"]["MySubTaskSet"]["tasks"]["task1"]["ratio"])
     self.assertEqual(0.125, ratio["MyTaskSet"]["tasks"]["MySubTaskSet"]["tasks"]["task2"]["ratio"])
Example #3
0
 def test_get_task_ratio_dict_relative(self):
     ratio = get_task_ratio_dict([MyTaskSet])
     self.assertEqual(1.0, ratio["MyTaskSet"]["ratio"])
     self.assertEqual(0.75, ratio["MyTaskSet"]["tasks"]["root_task"]["ratio"])
     self.assertEqual(0.25, ratio["MyTaskSet"]["tasks"]["MySubTaskSet"]["ratio"])
     self.assertEqual(0.5, ratio["MyTaskSet"]["tasks"]["MySubTaskSet"]["tasks"]["task1"]["ratio"])
     self.assertEqual(0.5, ratio["MyTaskSet"]["tasks"]["MySubTaskSet"]["tasks"]["task2"]["ratio"])
Example #4
0
 def test_get_task_ratio_dict_relative(self):
     ratio = get_task_ratio_dict([MyTaskSet])
     assert 1.0 == ratio["MyTaskSet"]["ratio"]
     assert 0.75 == ratio["MyTaskSet"]["tasks"]["root_task"]["ratio"]
     assert 0.25 == ratio["MyTaskSet"]["tasks"]["MySubTaskSet"]["ratio"]
     assert 0.5 == ratio["MyTaskSet"]["tasks"]["MySubTaskSet"]["tasks"]["task1"]["ratio"]
     assert 0.5 == ratio["MyTaskSet"]["tasks"]["MySubTaskSet"]["tasks"]["task2"]["ratio"]
Example #5
0
 def test_get_task_ratio_dict_relative(self):
     ratio = get_task_ratio_dict([MyTaskSet])
     self.assertEqual(1.0, ratio["MyTaskSet"]["ratio"])
     self.assertEqual(0.75, ratio["MyTaskSet"]["tasks"]["root_task"]["ratio"])
     self.assertEqual(0.25, ratio["MyTaskSet"]["tasks"]["MySubTaskSet"]["ratio"])
     self.assertEqual(0.5, ratio["MyTaskSet"]["tasks"]["MySubTaskSet"]["tasks"]["task1"]["ratio"])
     self.assertEqual(0.5, ratio["MyTaskSet"]["tasks"]["MySubTaskSet"]["tasks"]["task2"]["ratio"])
Example #6
0
 def test_get_task_ratio_dict_total(self):
     ratio = get_task_ratio_dict([MyTaskSet], total=True)
     assert 1.0 == ratio["MyTaskSet"]["ratio"]
     assert 0.75 == ratio["MyTaskSet"]["tasks"]["root_task"]["ratio"]
     assert 0.25 == ratio["MyTaskSet"]["tasks"]["MySubTaskSet"]["ratio"]
     assert 0.125 == ratio["MyTaskSet"]["tasks"]["MySubTaskSet"]["tasks"]["task1"]["ratio"]
     assert 0.125 == ratio["MyTaskSet"]["tasks"]["MySubTaskSet"]["tasks"]["task2"]["ratio"]
Example #7
0
    def test_task_ratio_command_with_locust_weight(self):
        class Tasks(TaskSet):
            @task(1)
            def task1(self):
                pass

            @task(3)
            def task3(self):
                pass

        class UnlikelyUser(User):
            weight = 1
            tasks = [Tasks]

        class MoreLikelyUser(User):
            weight = 3
            tasks = [Tasks]

        ratio_dict = get_task_ratio_dict([UnlikelyUser, MoreLikelyUser],
                                         total=True)

        self.assertDictEqual(
            {
                'UnlikelyUser': {
                    'ratio': 0.25,
                    'tasks': {
                        'Tasks': {
                            'tasks': {
                                'task1': {
                                    'ratio': 0.25 * 0.25
                                },
                                'task3': {
                                    'ratio': 0.25 * 0.75
                                },
                            },
                            'ratio': 0.25
                        }
                    },
                },
                'MoreLikelyUser': {
                    'ratio': 0.75,
                    'tasks': {
                        'Tasks': {
                            'tasks': {
                                'task1': {
                                    'ratio': 0.75 * 0.25
                                },
                                'task3': {
                                    'ratio': 0.75 * 0.75
                                },
                            },
                            'ratio': 0.75,
                        },
                    },
                }
            }, ratio_dict)
        unlikely = ratio_dict['UnlikelyUser']['tasks']['Tasks']['tasks']
        likely = ratio_dict['MoreLikelyUser']['tasks']['Tasks']['tasks']
        assert unlikely['task1']['ratio'] + unlikely['task3'][
            'ratio'] + likely['task1']['ratio'] + likely['task3']['ratio'] == 1
Example #8
0
    def test_task_ratio_command_with_locust_weight(self):
        class Tasks(TaskSet):
            @task(1)
            def task1(self):
                pass

            @task(3)
            def task3(self):
                pass

        class UnlikelyLocust(Locust):
            weight = 1
            task_set = Tasks

        class MoreLikelyLocust(Locust):
            weight = 3
            task_set = Tasks

        ratio_dict = get_task_ratio_dict([UnlikelyLocust, MoreLikelyLocust], total=True)

        self.assertEqual({
               'UnlikelyLocust':   {'tasks': {'task1': {'ratio': 0.25*0.25}, 'task3': {'ratio': 0.25*0.75}}, 'ratio': 0.25},
               'MoreLikelyLocust': {'tasks': {'task1': {'ratio': 0.75*0.25}, 'task3': {'ratio': 0.75*0.75}}, 'ratio': 0.75}
               }, ratio_dict)
        unlikely = ratio_dict['UnlikelyLocust']['tasks']
        likely = ratio_dict['MoreLikelyLocust']['tasks']
        assert unlikely['task1']['ratio'] + unlikely['task3']['ratio'] + likely['task1']['ratio'] + likely['task3']['ratio'] == 1
Example #9
0
 def test_task_ratio_command(self):
     class Tasks(TaskSet):
         @task
         def root_task1(self):
             pass
         
         @task
         class SubTasks(TaskSet):
             @task
             def task1(self):
                 pass
             
             @task
             def task2(self):
                 pass
     
     class User(Locust):
         task_set = Tasks
     
     ratio_dict = get_task_ratio_dict(User.task_set.tasks, total=True)
     
     self.assertEqual({
         'SubTasks': {
             'tasks': {
                 'task1': {'ratio': 0.25},
                 'task2': {'ratio': 0.25}
             },
             'ratio': 0.5
         }, 
         'root_task1': {'ratio': 0.5}
     }, ratio_dict)
Example #10
0
 def test_task_ratio_command(self):
     class Tasks(TaskSet):
         @task
         def root_task1(self):
             pass
         
         @task
         class SubTasks(TaskSet):
             @task
             def task1(self):
                 pass
             
             @task
             def task2(self):
                 pass
     
     class User(Locust):
         task_set = Tasks
     
     ratio_dict = get_task_ratio_dict(User.task_set.tasks, total=True)
     
     self.assertEqual({
         'SubTasks': {
             'tasks': {
                 'task1': {'ratio': 0.25},
                 'task2': {'ratio': 0.25}
             },
             'ratio': 0.5
         }, 
         'root_task1': {'ratio': 0.5}
     }, ratio_dict)
Example #11
0
    def test_task_ratio_command_with_locust_weight(self):
        class Tasks(TaskSet):
            @task(1)
            def task1(self):
                pass

            @task(3)
            def task3(self):
                pass

        class UnlikelyLocust(Locust):
            weight = 1
            task_set = Tasks

        class MoreLikelyLocust(Locust):
            weight = 3
            task_set = Tasks

        ratio_dict = get_task_ratio_dict([UnlikelyLocust, MoreLikelyLocust], total=True)

        self.assertEqual({
               'UnlikelyLocust':   {'tasks': {'task1': {'ratio': 0.25*0.25}, 'task3': {'ratio': 0.25*0.75}}, 'ratio': 0.25},
               'MoreLikelyLocust': {'tasks': {'task1': {'ratio': 0.75*0.25}, 'task3': {'ratio': 0.75*0.75}}, 'ratio': 0.75}
               }, ratio_dict)
        unlikely = ratio_dict['UnlikelyLocust']['tasks']
        likely = ratio_dict['MoreLikelyLocust']['tasks']
        assert unlikely['task1']['ratio'] + unlikely['task3']['ratio'] + likely['task1']['ratio'] + likely['task3']['ratio'] == 1
Example #12
0
    def test_task_ratio_command_with_locust_weight(self):
        class Tasks(TaskSet):
            @task(1)
            def task1(self):
                pass

            @task(3)
            def task3(self):
                pass

        class UnlikelyLocust(Locust):
            weight = 1
            task_set = Tasks

        class MoreLikelyLocust(Locust):
            weight = 3
            task_set = Tasks

        ratio_dict = get_task_ratio_dict([UnlikelyLocust, MoreLikelyLocust],
                                         total=True)

        self.assertEqual(
            {
                "UnlikelyLocust": {
                    "tasks": {
                        "task1": {
                            "ratio": 0.25 * 0.25
                        },
                        "task3": {
                            "ratio": 0.25 * 0.75
                        },
                    },
                    "ratio": 0.25,
                },
                "MoreLikelyLocust": {
                    "tasks": {
                        "task1": {
                            "ratio": 0.75 * 0.25
                        },
                        "task3": {
                            "ratio": 0.75 * 0.75
                        },
                    },
                    "ratio": 0.75,
                },
            },
            ratio_dict,
        )
        unlikely = ratio_dict["UnlikelyLocust"]["tasks"]
        likely = ratio_dict["MoreLikelyLocust"]["tasks"]
        assert (unlikely["task1"]["ratio"] + unlikely["task3"]["ratio"] +
                likely["task1"]["ratio"] + likely["task3"]["ratio"] == 1)
Example #13
0
    def test_task_ratio_command(self):
        class Tasks(TaskSet):
            @task
            def root_task1(self):
                pass

            @task
            class SubTasks(TaskSet):
                @task
                def task1(self):
                    pass

                @task
                def task2(self):
                    pass

        class User(Locust):
            task_set = Tasks

        ratio_dict = get_task_ratio_dict(User.task_set.tasks, total=True)

        self.assertEqual(
            {
                "SubTasks": {
                    "tasks": {
                        "task1": {
                            "ratio": 0.25
                        },
                        "task2": {
                            "ratio": 0.25
                        }
                    },
                    "ratio": 0.5,
                },
                "root_task1": {
                    "ratio": 0.5
                },
            },
            ratio_dict,
        )
Example #14
0
def main():
    parser, options, arguments = parse_options()

    # setup logging
    
    if options.show_version:
        print("Locust %s" % (version,))
        sys.exit(0)

    locustfile = find_locustfile(options.locustfile)
    if not locustfile:
        logger.error("Could not find any locustfile! Ensure file ends in '.py' and see --help for available options.")
        sys.exit(1)

    docstring, locusts = load_locustfile(locustfile)

    if options.list_commands:
        logger.info("Available Locusts:")
        for name in locusts:
            logger.info("    " + name)
        sys.exit(0)

    if not locusts:
        logger.error("No Locust class found!")
        sys.exit(1)

    # make sure specified Locust exists
    if arguments:
        missing = set(arguments) - set(locusts.keys())
        if missing:
            logger.error("Unknown Locust(s): %s\n" % (", ".join(missing)))
            sys.exit(1)
        else:
            names = set(arguments) & set(locusts.keys())
            locust_classes = [locusts[n] for n in names]
    else:
        locust_classes = locusts.values()
    
    if options.show_task_ratio:
        logger.info("\n Task ratio per locust class")
        logger.info( "-" * 80)
        print_task_ratio(locust_classes)
        logger.info("\n Total task ratio")
        logger.info("-" * 80)
        print_task_ratio(locust_classes, total=True)
        sys.exit(0)
    if options.show_task_ratio_json:
        from json import dumps
        task_data = {
            "per_class": get_task_ratio_dict(locust_classes), 
            "total": get_task_ratio_dict(locust_classes, total=True)
        }
        logger.info(dumps(task_data))
        sys.exit(0)
    
    # if --master is set, make sure --no-web isn't set
    if options.master and options.no_web:
        logger.error("Locust can not run distributed with the web interface disabled (do not use --no-web and --master together)")
        sys.exit(0)

    if not options.no_web and not options.slave:
        # spawn web greenlet
        logger.info("Starting web monitor at %s:%s" % (options.web_host or "*", options.port))
        main_greenlet = gevent.spawn(web.start, locust_classes, options)
    
    if not options.master and not options.slave:
        runners.locust_runner = LocalLocustRunner(locust_classes, options)
        # spawn client spawning/hatching greenlet
        if options.no_web:
            runners.locust_runner.start_hatching(wait=True)
            main_greenlet = runners.locust_runner.greenlet
    elif options.master:
        runners.locust_runner = MasterLocustRunner(locust_classes, options)
    elif options.slave:
        try:
            runners.locust_runner = SlaveLocustRunner(locust_classes, options)
            main_greenlet = runners.locust_runner.greenlet
        except socket.error as e:
            logger.error("Failed to connect to the Locust master: %s", e)
            sys.exit(-1)
    
    if not options.only_summary and (options.print_stats or (options.no_web and not options.slave)):
        # spawn stats printing greenlet
        gevent.spawn(stats_printer)
    
    def shutdown(code=0):
        """
        Shut down locust by firing quitting event, printing stats and exiting
        """
        logger.info("Shutting down (exit code %s), bye." % code)

        events.quitting.fire()
        print_stats(runners.locust_runner.request_stats)
        print_percentile_stats(runners.locust_runner.request_stats)

        print_error_report()
        sys.exit(code)
    
    # install SIGTERM handler
    def sig_term_handler():
        logger.info("Got SIGTERM signal")
        shutdown(0)
    gevent.signal(signal.SIGTERM, sig_term_handler)
    
    try:
        logger.info("Starting Locust %s" % version)
        main_greenlet.join()
        code = 0
        if len(runners.locust_runner.errors):
            code = 1
        shutdown(code=code)
    except KeyboardInterrupt as e:
        shutdown(0)
Example #15
0
def main():
    parser, options = parse_options()

    # setup logging
    if not options.skip_log_setup:
        setup_logging(options.loglevel, options.logfile)

    logger = logging.getLogger(__name__)

    locust_path = get_locust_path()

    if options.web:
        web_port = options.web_port
        logger.info('Running easy-locust web: 0.0.0.0:{}'.format(web_port))
        init_app(port=web_port)
        sys.exit(0)

    if options.demo:
        if not locust_path:
            logger.error(
                '''Cannot locate Python path, make sure it is in right place. If windows add it to sys PATH,
            if linux make sure python is installed in /usr/local/lib/''')
            sys.exit(1)
        pt_demo_path = os.path.join(locust_path, 'demo',
                                    'demo_pressuretest.xls')
        pt_demo_path_json = os.path.join(locust_path, 'demo',
                                         'demo_locustfile.json')
        pt_new_demo = os.path.join(os.getcwd(), 'PtDemo.xls')
        pt_new_demo_json = os.path.join(os.getcwd(), 'demo.json')
        shutil.copyfile(pt_demo_path, pt_new_demo)
        shutil.copyfile(pt_demo_path_json, pt_new_demo_json)
        sys.exit(0)

    if options.xlsfile:
        pt_file = options.xlsfile
        if not (pt_file.endswith('.xls') or pt_file.endswith('.json')):
            logger.error(
                "PressureTest file must be end with '.xls' or '.json' and see --help for available options."
            )
            sys.exit(1)
        if not os.path.isfile(pt_file):
            logger.error('PressureTest file is not exist, please check it.')
            sys.exit(1)
        _status = generate_locust_file(pt_file)
        if not _status:
            sys.exit(1)
        sys.exit(0)

    locustfile = find_locustfile(options.locustfile)

    if not locustfile:
        logger.error(
            "Could not find any locustfile! Ensure file ends in '.py' and see --help for available options."
        )
        sys.exit(1)

    if locustfile == "locust.py" or locustfile == "locust.xls" or locustfile == "locust.json":
        logger.error(
            "The locustfile must not be named `locust.py` or `locust.xls` or `locust.json`. "
            "Please rename the file and try again.")
        sys.exit(1)

    docstring, locusts = load_locustfile(locustfile)

    if options.list_commands:
        console_logger.info("Available Locusts:")
        for name in locusts:
            console_logger.info("    " + name)
        sys.exit(0)

    if not locusts:
        logger.error("No Locust class found!")
        sys.exit(1)

    # make sure specified Locust exists
    if options.locust_classes:
        missing = set(options.locust_classes) - set(locusts.keys())
        if missing:
            logger.error("Unknown Locust(s): %s\n" % (", ".join(missing)))
            sys.exit(1)
        else:
            names = set(options.locust_classes) & set(locusts.keys())
            locust_classes = [locusts[n] for n in names]
    else:
        # list() call is needed to consume the dict_view object in Python 3
        locust_classes = list(locusts.values())

    if options.show_task_ratio:
        console_logger.info("\n Task ratio per locust class")
        console_logger.info("-" * 80)
        print_task_ratio(locust_classes)
        console_logger.info("\n Total task ratio")
        console_logger.info("-" * 80)
        print_task_ratio(locust_classes, total=True)
        sys.exit(0)
    if options.show_task_ratio_json:
        from json import dumps
        task_data = {
            "per_class": get_task_ratio_dict(locust_classes),
            "total": get_task_ratio_dict(locust_classes, total=True)
        }
        console_logger.info(dumps(task_data))
        sys.exit(0)

    if options.run_time:
        if not options.no_web:
            logger.error(
                "The --run-time argument can only be used together with --no-web"
            )
            sys.exit(1)
        try:
            options.run_time = parse_timespan(options.run_time)
        except ValueError:
            logger.error(
                "Valid --run-time formats are: 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc."
            )
            sys.exit(1)

        def spawn_run_time_limit_greenlet():
            logger.info("Run time limit set to %s seconds" % options.run_time)

            def timelimit_stop():
                logger.info("Time limit reached. Stopping Locust.")
                runners.locust_runner.quit()

            gevent.spawn_later(options.run_time, timelimit_stop)

    if options.step_time:
        if not options.step_load:
            logger.error(
                "The --step-time argument can only be used together with --step-load"
            )
            sys.exit(1)
        try:
            options.step_time = parse_timespan(options.step_time)
        except ValueError:
            logger.error(
                "Valid --step-time formats are: 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc."
            )
            sys.exit(1)

    if options.master:
        # Add -d for automatically run slaves
        if options.distribute:
            ptpy = locustfile
            if options.locustfile.endswith('.xls'):
                _type = 'xls'
                pt_s = PtExcel(options.locustfile)
                master_ip, pt_slave_info = pt_s.pt_slave()
            else:
                _type = 'dict'
                with open(options.locustfile, 'r') as f:
                    _d = json.load(f, encoding='utf-8')
                master_ip = _d.get('master_ip')
                pt_slave_info = _d.get('slaves')
            if master_ip == '':
                logger.error(
                    'master IP cannot be None if you use --distribute')
                sys.exit(1)
            if options.boomer:
                locust_cli_slave = 'nohup ./client_v1 --web --master-host={masteIP} > /dev/null 2>&1 &'.format(
                    masteIP=master_ip)
                targets_dict, file_list = gen_boomer_client_json(
                    options.locustfile)
                boomer_client_file = os.path.join(locust_path, 'boomer_client',
                                                  'client_v1')
                file_list.append(boomer_client_file)
                thread_pool = []
                try:
                    for slave in pt_slave_info:
                        if _type == 'xls':
                            slave_ip, slave_username, slave_password = slave
                        else:
                            slave_ip, slave_username, slave_password = slave[
                                'ip'], slave['username'], slave['password']
                        _t = Thread(target=pt_slave_boomer,
                                    args=(slave_ip, slave_username,
                                          slave_password, file_list,
                                          locust_cli_slave, targets_dict))
                        logger.info('Prepare slave {}'.format(slave_ip))
                        thread_pool.append(_t)
                        _t.start()
                    for each_t in thread_pool:
                        each_t.join()
                    file_list.pop()
                    for each in file_list:
                        os.remove(each)
                except KeyboardInterrupt:
                    pass
                except Exception as e:
                    logger.error(
                        'Something happened, collect Exceptions here: {}'.
                        format(e))
            else:
                try:
                    locust_cli_slave = 'nohup locust -f /root/locust_client.py --slave --master-host={masteIP} > /dev/null 2>&1 &'\
                        .format(masteIP=master_ip)
                    thread_pool = []
                    for slave in pt_slave_info:
                        if _type == 'xls':
                            slave_ip, slave_username, slave_password = slave
                        else:
                            slave_ip, slave_username, slave_password = slave[
                                'ip'], slave['username'], slave['password']
                        _t = Thread(target=pt_slave,
                                    args=(slave_ip, slave_username,
                                          slave_password, ptpy,
                                          locust_cli_slave))
                        logger.info('Prepare slave {}'.format(slave_ip))
                        thread_pool.append(_t)
                        _t.start()
                    for each_t in thread_pool:
                        each_t.join()
                except KeyboardInterrupt:
                    pass
                except Exception as e:
                    logger.error(
                        'Something happened, collect Exceptions here: {}'.
                        format(e))

        runners.locust_runner = MasterLocustRunner(locust_classes, options)
    else:
        runners.locust_runner = LocalLocustRunner(locust_classes, options)
    # main_greenlet is pointing to runners.locust_runner.greenlet by default, it will point the web greenlet later if in web mode
    main_greenlet = runners.locust_runner.greenlet

    if options.no_web:
        if options.master:
            while len(runners.locust_runner.clients.ready
                      ) < options.expect_slaves:
                logging.info(
                    "Waiting for slaves to be ready, %s of %s connected",
                    len(runners.locust_runner.clients.ready),
                    options.expect_slaves)
                time.sleep(1)
        if options.step_time:
            runners.locust_runner.start_stepload(options.num_clients,
                                                 options.hatch_rate,
                                                 options.step_clients,
                                                 options.step_time)
        else:
            runners.locust_runner.start_hatching(options.num_clients,
                                                 options.hatch_rate)
            # make locusts are spawned
            time.sleep(1)
    else:
        # spawn web greenlet
        logger.info("Starting web monitor at http://%s:%s" %
                    (options.web_host or "*", options.port))
        main_greenlet = gevent.spawn(web.start, locust_classes, options)

    if options.run_time:
        spawn_run_time_limit_greenlet()

    stats_printer_greenlet = None
    if not options.only_summary and (options.print_stats or options.no_web):
        # spawn stats printing greenlet
        stats_printer_greenlet = gevent.spawn(stats_printer)

    if options.csvfilebase:
        gevent.spawn(stats_writer, options.csvfilebase,
                     options.stats_history_enabled)

    def shutdown(code=0):
        """
        Shut down locust by firing quitting event, printing/writing stats and exiting
        """
        logger.info("Shutting down (exit code %s), bye." % code)
        if stats_printer_greenlet is not None:
            stats_printer_greenlet.kill(block=False)
        logger.info("Cleaning up runner...")
        if runners.locust_runner is not None:
            runners.locust_runner.quit()
        logger.info("Running teardowns...")
        events.quitting.fire(reverse=True)
        print_stats(runners.locust_runner.stats, current=False)
        print_percentile_stats(runners.locust_runner.stats)
        if options.csvfilebase:
            write_stat_csvs(options.csvfilebase, options.stats_history_enabled)
        print_error_report()
        sys.exit(code)

    # install SIGTERM handler
    def sig_term_handler():
        logger.info("Got SIGTERM signal")
        shutdown(0)

    gevent.signal(signal.SIGTERM, sig_term_handler)

    try:
        logger.info("Starting Locust %s" % version)
        main_greenlet.join()
        code = 0
        if len(runners.locust_runner.errors) or len(
                runners.locust_runner.exceptions):
            code = options.exit_code_on_error
        shutdown(code=code)
    except KeyboardInterrupt:
        shutdown(0)
Example #16
0
def main(user_params=None, common_param=None):
    parser, options, arguments = parse_options()

    # setup logging
    setup_logging(options.loglevel, options.logfile)
    logger = logging.getLogger(__name__)

    if options.show_version:
        print("Locust %s" % (version, ))
        sys.exit(0)

    locustfile = find_locustfile(options.locustfile)

    if not locustfile:
        logger.error(
            "Could not find any locustfile! Ensure file ends in '.py' and see --help for available options."
        )
        sys.exit(1)

    if locustfile == "locust.py":
        logger.error(
            "The locustfile must not be named `locust.py`. Please rename the file and try again."
        )
        sys.exit(1)

    docstring, locusts = load_locustfile(locustfile)

    if options.list_commands:
        console_logger.info("Available Locusts:")
        for name in locusts:
            console_logger.info("    " + name)
        sys.exit(0)

    if not locusts:
        logger.error("No Locust class found!")
        sys.exit(1)

    # make sure specified Locust exists
    if arguments:
        missing = set(arguments) - set(locusts.keys())
        if missing:
            logger.error("Unknown Locust(s): %s\n" % (", ".join(missing)))
            sys.exit(1)
        else:
            names = set(arguments) & set(locusts.keys())
            locust_classes = [locusts[n] for n in names]
    else:
        # list() call is needed to consume the dict_view object in Python 3
        locust_classes = list(locusts.values())

    if options.show_task_ratio:
        console_logger.info("\n Task ratio per locust class")
        console_logger.info("-" * 80)
        print_task_ratio(locust_classes)
        console_logger.info("\n Total task ratio")
        console_logger.info("-" * 80)
        print_task_ratio(locust_classes, total=True)
        sys.exit(0)
    if options.show_task_ratio_json:
        from json import dumps
        task_data = {
            "per_class": get_task_ratio_dict(locust_classes),
            "total": get_task_ratio_dict(locust_classes, total=True)
        }
        console_logger.info(dumps(task_data))
        sys.exit(0)

    if not options.no_web and not options.slave:
        # spawn web greenlet
        logger.info("Starting web monitor at %s:%s" %
                    (options.web_host or "*", options.port))
        main_greenlet = gevent.spawn(web.start, locust_classes, options)

    if not options.master and not options.slave:
        if user_params:
            runners.locust_runner = ParameterizableLocustRunner(
                locust_classes, options, user_params, common_param)
        else:
            runners.locust_runner = LocalLocustRunner(locust_classes, options)
        # spawn client spawning/hatching greenlet
        if options.no_web:
            runners.locust_runner.start_hatching(wait=True)
            main_greenlet = runners.locust_runner.greenlet
    elif options.master:
        runners.locust_runner = MasterLocustRunner(locust_classes, options)
        if options.no_web:
            while len(runners.locust_runner.clients.ready
                      ) < options.expect_slaves:
                logging.info(
                    "Waiting for slaves to be ready, %s of %s connected",
                    len(runners.locust_runner.clients.ready),
                    options.expect_slaves)
                time.sleep(1)

            runners.locust_runner.start_hatching(options.num_clients,
                                                 options.hatch_rate)
            main_greenlet = runners.locust_runner.greenlet
    elif options.slave:
        try:
            runners.locust_runner = SlaveLocustRunner(locust_classes, options)
            main_greenlet = runners.locust_runner.greenlet
        except socket.error as e:
            logger.error("Failed to connect to the Locust master: %s", e)
            sys.exit(-1)

    if not options.only_summary and (options.print_stats or
                                     (options.no_web and not options.slave)):
        # spawn stats printing greenlet
        gevent.spawn(stats_printer)

    if options.csvfilebase:
        gevent.spawn(stats_writer, options.csvfilebase)

    def shutdown(code=0):
        """
        Shut down locust by firing quitting event, printing/writing stats and exiting
        """
        logger.info("Shutting down (exit code %s), bye." % code)

        events.quitting.fire()
        print_stats(runners.locust_runner.request_stats)
        print_percentile_stats(runners.locust_runner.request_stats)
        if options.csvfilebase:
            write_stat_csvs(options.csvfilebase)
        print_error_report()
        sys.exit(code)

    # install SIGTERM handler
    def sig_term_handler():
        logger.info("Got SIGTERM signal")
        shutdown(0)

    gevent.signal(gevent.signal.SIGTERM, sig_term_handler)

    try:
        logger.info("Starting Locust %s" % version)
        main_greenlet.join()
        code = 0
        if len(runners.locust_runner.errors):
            code = 1
        shutdown(code=code)
    except KeyboardInterrupt as e:
        shutdown(0)