Merge pull request locustio#38 from erlanggakrisnamukti/develop

Switch File Feature to Master

locust/main.py

@@ -1,6 +1,7 @@
 import locust
 from . import runners
 
+import imp
 import gevent
 import sys
 import os
@@ -45,15 +46,15 @@ def parse_options():
         default="",
         help="Host to bind the web interface to. Defaults to '' (all interfaces)"
     )
-    
+
     parser.add_option(
         '-P', '--port', '--web-port',
         type="int",
         dest="port",
         default=8089,
         help="Port on which to run web host"
     )
-    
+
     parser.add_option(
         '-f', '--locustfile',
         dest='locustfile',
@@ -78,7 +79,7 @@ def parse_options():
         default=False,
         help="Set locust to run in distributed mode with this process as slave"
     )
-    
+
     # master host options
     parser.add_option(
         '--master-host',
@@ -88,7 +89,7 @@ def parse_options():
         default="127.0.0.1",
         help="Host or IP address of locust master for distributed load testing. Only used when running with --slave. Defaults to 127.0.0.1."
     )
-    
+
     parser.add_option(
         '--master-port',
         action='store',
@@ -106,7 +107,7 @@ def parse_options():
         default="*",
         help="Interfaces (hostname, ip) that locust master should bind to. Only used when running with --master. Defaults to * (all available interfaces)."
     )
-    
+
     parser.add_option(
         '--master-bind-port',
         action='store',
@@ -153,7 +154,7 @@ def parse_options():
         default=1,
         help="The rate per second in which clients are spawned. Only used together with --no-web"
     )
-    
+
     # Number of requests
     parser.add_option(
         '-n', '--num-request',
@@ -163,7 +164,7 @@ def parse_options():
         default=None,
         help="Number of requests to perform. Only used together with --no-web"
     )
-    
+
     # log level
     parser.add_option(
         '--loglevel', '-L',
@@ -173,7 +174,7 @@ def parse_options():
         default='INFO',
         help="Choose between DEBUG/INFO/WARNING/ERROR/CRITICAL. Default is INFO.",
     )
-    
+
     # log file
     parser.add_option(
         '--logfile',
@@ -183,7 +184,7 @@ def parse_options():
         default=None,
         help="Path to log file. If not set, log will go to stdout/stderr",
     )
-    
+
     # if we should print stats in the console
     parser.add_option(
         '--print-stats',
@@ -209,7 +210,7 @@ def parse_options():
         default=False,
         help="Do not reset statistics once hatching has been completed",
     )
-    
+
     # List locust commands found in loaded locust files/source files
     parser.add_option(
         '-l', '--list',
@@ -218,7 +219,7 @@ def parse_options():
         default=False,
         help="Show list of possible locust classes and exit"
     )
-    
+
     # Display ratio table of all tasks
     parser.add_option(
         '--show-task-ratio',
@@ -235,7 +236,7 @@ def parse_options():
         default=False,
         help="print json data of the locust classes' task execution ratio"
     )
-    
+
     # Version number (optparse gives you --version but we have to do it
     # ourselves to get -V too. sigh)
     parser.add_option(
@@ -316,6 +317,23 @@ def is_locust(tup):
         and not name.startswith('_')
     )
 
+def truncate_path(path):
+    # split path which comes from command on terminal
+    splitted_path = os.path.normpath(path).split(os.path.sep)
+
+    count = 0
+    for i in reversed(xrange(len(splitted_path))):
+        if count < 3 and splitted_path[i]:
+            if count == 0:
+                final_path = splitted_path[i]
+            elif count == 2:
+                final_path = os.path.join("...", splitted_path[i], final_path)
+            else:
+                final_path = os.path.join(splitted_path[i], final_path)
+            count += 1
+        else:
+            break
+    return final_path
 
 def load_locustfile(path):
     """
@@ -345,7 +363,7 @@ def load_locustfile(path):
             sys.path.insert(0, directory)
             del sys.path[i + 1]
     # Perform the import (trimming off the .py)
-    imported = __import__(os.path.splitext(locustfile)[0])
+    imported = imp.load_source(os.path.splitext(locustfile)[0], path)
     # Remove directory from path if we added it ourselves (just to be neat)
     if added_to_path:
         del sys.path[0]
@@ -355,30 +373,54 @@ def load_locustfile(path):
         del sys.path[0]
     # Return our two-tuple
     locusts = dict(filter(is_locust, vars(imported).items()))
-    return imported.__doc__, locusts
+
+    # truncate the fullpath
+    final_path = truncate_path(path)
+
+    return {final_path: locusts}
+
+def collect_locustfiles(path):
+    collected = dict()
+
+    for root, dirs, files in os.walk(path):
+        if files:
+            for file_ in files:
+                if file_.endswith('.py') and not file_.endswith('__init__.py'):
+                    fullpath = os.path.abspath(os.path.join(root, file_))
+                    loaded = load_locustfile(fullpath)
+                    if loaded:
+                        collected.update(loaded)
+    return collected
 
 def main():
     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 os.path.isdir(options.locustfile):
+        all_locustfiles = collect_locustfiles(options.locustfile)
+    else:
+        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 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)
+        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)
+        # docstring, all_locustsfiles = load_locustfile(locustfile)
+        all_locustfiles = load_locustfile(locustfile)
+
+    # Use the first locustfile for the default locusts
+    locusts = all_locustfiles.values()[0]
 
     if options.list_commands:
         console_logger.info("Available Locusts:")
@@ -402,7 +444,7 @@ def main():
     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)
@@ -414,7 +456,7 @@ def main():
     if options.show_task_ratio_json:
         from json import dumps
         task_data = {
-            "per_class": get_task_ratio_dict(locust_classes), 
+            "per_class": get_task_ratio_dict(locust_classes),
             "total": get_task_ratio_dict(locust_classes, total=True)
         }
         console_logger.info(dumps(task_data))
@@ -424,15 +466,15 @@ def main():
         # 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)
+        runners.locust_runner = LocalLocustRunner(locust_classes, options, available_locustfiles=all_locustfiles)
         # 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)
+        runners.locust_runner = MasterLocustRunner(locust_classes, options, available_locustfiles=all_locustfiles)
         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",
@@ -443,16 +485,16 @@ def main():
             main_greenlet = runners.locust_runner.greenlet
     elif options.slave:
         try:
-            runners.locust_runner = SlaveLocustRunner(locust_classes, options)
+            runners.locust_runner = SlaveLocustRunner(locust_classes, options, available_locustfiles=all_locustfiles)
             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
@@ -465,13 +507,13 @@ def shutdown(code=0):
 
         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()

locust/runners.py

@@ -27,11 +27,11 @@
 SLAVE_REPORT_INTERVAL = 3.0
 NORMAL, RAMP = ["Normal", "Auto"]
 
-
 class LocustRunner(object):
-    def __init__(self, locust_classes, options):
+    def __init__(self, locust_classes, options, available_locustfiles=None):
         self.options = options
         self.locust_classes = locust_classes
+        self.available_locustfiles = available_locustfiles or {}
         self.hatch_rate = options.hatch_rate
         self.num_clients = options.num_clients
         self.num_requests = options.num_requests
@@ -148,6 +148,16 @@ def kill_locusts(self, kill_count):
             self.locusts.killone(g)
         events.hatch_complete.fire(user_count=self.num_clients)
 
+    def select_file(self, key):
+        """
+            Set the active locust classes to the executeables described by the key
+        """
+        try:
+            self.locust_classes = self.available_locustfiles[key].values()
+        except KeyError:
+            logger.error("No available locust classes found with key: {}".format(key))
+            self.locust_classes = []
+
     def start_hatching(self, locust_count=None, hatch_rate=None, wait=False):
         if self.state != STATE_RUNNING and self.state != STATE_HATCHING:
             self.stats.clear_all()
@@ -194,9 +204,8 @@ def log_exception(self, node_id, msg, formatted_tb):
         self.exceptions[key] = row
 
 class LocalLocustRunner(LocustRunner):
-    def __init__(self, locust_classes, options):
-        super(LocalLocustRunner, self).__init__(locust_classes, options)
-
+    def __init__(self, locust_classes, options, available_locustfiles=None):
+        super(LocalLocustRunner, self).__init__(locust_classes, options, available_locustfiles)
         # register listener thats logs the exception for the local runner
         def on_locust_error(locust_instance, exception, tb):
             formatted_tb = "".join(traceback.format_tb(tb))
@@ -208,8 +217,8 @@ def start_hatching(self, locust_count=None, hatch_rate=None, wait=False):
         self.greenlet = self.hatching_greenlet
 
 class DistributedLocustRunner(LocustRunner):
-    def __init__(self, locust_classes, options):
-        super(DistributedLocustRunner, self).__init__(locust_classes, options)
+    def __init__(self, locust_classes, options, available_locustfiles=None):
+        super(DistributedLocustRunner, self).__init__(locust_classes, options, available_locustfiles)
         self.master_host = options.master_host
         self.master_port = options.master_port
         self.master_bind_host = options.master_bind_host

locust/static/locust.js

@@ -1,4 +1,5 @@
 $(window).ready(function() {
+    $('.select2').select2({theme: 'bootstrap'});
     if($("#locust_count").length > 0) {
         $("#locust_count").focus().select();
     }
@@ -202,6 +203,7 @@ function updateStats() {
         $("#status_text").html(report.state);
         $("#userCount").html(report.user_count);
         $("#running_type").html(report.running_type);
+        $("#host_url").html(report.host)
 
         if (typeof report.slave_count !== "undefined")
             $("#slaveCount").html(report.slave_count)