Added functionality for printing the task execution ratio in JSON for…

…mat using the --show-task-ratio-json flag

locust/inspectlocust.py

@@ -24,3 +24,28 @@ def print_task_ratio(locusts, total=False, level=0, parent_ratio=1.0):
 			else:
 				print_task_ratio(locust.tasks, total, level+1)
 
+
+def get_task_ratio_dict(tasks, total=False, parent_ratio=1.0):
+	"""
+	Return a dict containing task execution ratio info
+	"""
+	ratio = {}
+	for task in tasks:
+		ratio.setdefault(task, 0)
+		ratio[task] += 1
+
+	# get percentage
+	ratio_percent = dict(map(lambda x: (x[0], float(x[1])/len(tasks) * parent_ratio), ratio.iteritems()))
+
+	task_dict = {}
+	for locust, ratio in ratio_percent.iteritems():
+		d = {"ratio":ratio}
+		if inspect.isclass(locust) and issubclass(locust, LocustBase):
+			if total:
+				d["tasks"] = get_task_ratio_dict(locust.tasks, total, ratio)
+			else:
+				d["tasks"] = get_task_ratio_dict(locust.tasks, total)
+
+		task_dict[locust.__name__] = d
+
+	return task_dict

locust/main.py

@@ -13,7 +13,7 @@
 import web
 from log import setup_logging, console_logger
 from stats import stats_printer, RequestStats, print_percentile_stats, print_error_report, print_stats
-from inspectlocust import print_task_ratio
+from inspectlocust import print_task_ratio, get_task_ratio_dict
 from core import Locust, WebLocust
 from runners import MasterLocustRunner, SlaveLocustRunner, LocalLocustRunner
 
@@ -175,6 +175,14 @@ def parse_options():
         default=False,
         help="print table of the locust classes' task execution ratio"
     )
+    # Display ratio table of all tasks in JSON format
+    parser.add_option(
+        '--show-task-ratio-json',
+        action='store_true',
+        dest='show_task_ratio_json',
+        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)
@@ -339,6 +347,14 @@ def main():
         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 --master is set, make sure --no-web isn't set
     if options.master and options.no_web:

locust/test/runtests.py

@@ -1,6 +1,6 @@
 import unittest
 from locust_class import TestLocustClass, TestWebLocustClass
-from test_stats import TestRequestStats, TestRequestStatsWithWebserver
+from test_stats import TestRequestStats, TestRequestStatsWithWebserver, TestInspectLocust
 
 if __name__ == '__main__':
 	unittest.main()

locust/test/test_stats.py

@@ -3,7 +3,8 @@
 
 from testcases import WebserverTestCase
 from locust.stats import RequestStats
-from locust.core import Locust
+from locust.core import Locust, SubLocust, task
+from locust.inspectlocust import get_task_ratio_dict
 
 class TestRequestStats(unittest.TestCase):
     def setUp(self):
@@ -89,3 +90,36 @@ class MyLocust(Locust):
         self.assertEqual(RequestStats.get("/ultra_fast").avg_content_length, len("This is an ultra fast response"))
         locust.client.get("/ultra_fast")
         self.assertEqual(RequestStats.get("/ultra_fast").avg_content_length, len("This is an ultra fast response"))
+
+
+class MyLocust(Locust):
+    @task(75)
+    def root_task(self):
+        pass
+
+    @task(25)
+    class MySubLocust(SubLocust):
+        @task
+        def task1(self):
+            pass
+        @task
+        def task2(self):
+            pass
+
+class TestInspectLocust(unittest.TestCase):
+    def test_get_task_ratio_dict_relative(self):
+        ratio = get_task_ratio_dict([MyLocust])
+        self.assertEqual(1.0, ratio["MyLocust"]["ratio"])
+        self.assertEqual(0.75, ratio["MyLocust"]["tasks"]["root_task"]["ratio"])
+        self.assertEqual(0.25, ratio["MyLocust"]["tasks"]["MySubLocust"]["ratio"])
+        self.assertEqual(0.5, ratio["MyLocust"]["tasks"]["MySubLocust"]["tasks"]["task1"]["ratio"])
+        self.assertEqual(0.5, ratio["MyLocust"]["tasks"]["MySubLocust"]["tasks"]["task2"]["ratio"])
+
+    def test_get_task_ratio_dict_total(self):
+        ratio = get_task_ratio_dict([MyLocust], total=True)
+        self.assertEqual(1.0, ratio["MyLocust"]["ratio"])
+        self.assertEqual(0.75, ratio["MyLocust"]["tasks"]["root_task"]["ratio"])
+        self.assertEqual(0.25, ratio["MyLocust"]["tasks"]["MySubLocust"]["ratio"])
+        self.assertEqual(0.125, ratio["MyLocust"]["tasks"]["MySubLocust"]["tasks"]["task1"]["ratio"])
+        self.assertEqual(0.125, ratio["MyLocust"]["tasks"]["MySubLocust"]["tasks"]["task2"]["ratio"])
+