Enhance (#35)

* enhance bench, can use mote k6 options

* fix typo

* fix github action

.github/workflows/import.yaml

@@ -84,4 +84,4 @@ jobs:
 
       - name: run stress testing
         run: |
-          python3 run.py stress run -d 3
+          python3 run.py stress run --args='-d 3s'

README.md

@@ -114,13 +114,27 @@ python3 run.py stress run --help
 python3 run.py stress run 
 
 # run all scenarios with 10 virtual users, every scenario lasts 3 seconds.
-python3 run.py stress run -vu 10 -d 3
+python3 run.py stress run --args='-u 10 -d 3s'
 
 # list all stress test scenarios
 python3 run.py stress scenarios
 
 # run go.Go1Step scenarios with 10 virtual users, every scenario lasts 3 seconds.
-python3 run.py stress run -vu 10 -d 3 -scenario go.Go1Step
+python3 run.py stress run -scenario go.Go1Step --args='-u 10 -d 3s'
+
+# run go.Go1Step scenarios with special test stage.
+# ramping up from 0 to 10 vus in first 10 seconds, then run 10 vus in 30 seconds, 
+# then ramping up from 10 to 50 vus in 10 seconds.
+python3 run.py stress run -scenario go.Go1Step --args='-s 10s:10 -s 30s:10 -s 10s:50'
+
+# use csv output
+python3 run.py stress run -scenario go.Go1Step --args='-s 10s:10 -s 30s:10 -s 10s:50 -o csv=test.csv'
+```
+
+for more k6 args, please refer to k6 run help.
+
+```bash
+scripts/k6 run --help
 ```
 
 k6 config file, summary result and outputs are in `output` folder. e.g.

README_cn.md

@@ -118,13 +118,27 @@ python3 run.py stress run --help
 python3 run.py stress run 
 
 # run all scenarios with 10 virtual users, every scenario lasts 3 seconds.
-python3 run.py stress run -vu 10 -d 3
+python3 run.py stress run --args='-u 10 -d 3s'
 
 # list all stress test scenarios
 python3 run.py stress scenarios
 
 # run go.Go1Step scenarios with 10 virtual users, every scenario lasts 3 seconds.
-python3 run.py stress run -vu 10 -d 3 -scenario go.Go1Step
+python3 run.py stress run -scenario go.Go1Step --args='-u 10 -d 3s'
+
+# run go.Go1Step scenarios with special test stage.
+# ramping up from 0 to 10 vus in first 10 seconds, then run 10 vus in 30 seconds, 
+# then ramping up from 10 to 50 vus in 10 seconds.
+python3 run.py stress run -scenario go.Go1Step --args='-s 10s:10 -s 30s:10 -s 10s:50'
+
+# use csv output
+python3 run.py stress run -scenario go.Go1Step --args='-s 10s:10 -s 30s:10 -s 10s:50 -o csv=test.csv'
+```
+
+更多 k6 参数，请参考。
+
+```bash
+scripts/k6 run --help
 ```
 
 k6 config file, summary result and outputs are in `output` folder. e.g.

nebula_bench/cli.py

@@ -132,10 +132,6 @@ def stress():
     default="int",
     help="space vid type, values should be [int, string], default: int",
 )
-@click.option("-vu", default=100, help="concurrent virtual users, default: 100")
-@click.option(
-    "-d", "--duration", default=60, help="duration for every scenario, unit: second, default: 60"
-)
 @click.option("-scenario", default="all", help="run special scenario, e.g. go.Go1Step")
 @click.option("-c", "--controller", default="k6", help="using which test tool")
 @click.option(
@@ -144,8 +140,9 @@ def stress():
     is_flag=True,
     help="Dry run, just dump stress testing config file, default: False",
 )
+@click.option("--args", help="extend args for test tool")
 def run(
-    folder, address, user, password, space, vid_type, scenario, controller, vu, duration, dry_run
+    folder, address, user, password, space, vid_type, scenario, controller, args, dry_run
 ):
     stress = StressFactory.gen_stress(
         _type=controller,
@@ -156,8 +153,7 @@ def run(
         space=space,
         vid_type=vid_type,
         scenarios=scenario,
-        vu=vu,
-        duration=duration,
+        args = args,
         dry_run=dry_run,
     )
     stress.run()

nebula_bench/stress.py

@@ -22,19 +22,11 @@ def load_scenarios(scenarios):
 
 
 class Stress(object):
+    DEFAULT_VU = 100
+    DEFAULT_DURATION = "60s"
+
     def __init__(
-        self,
-        folder,
-        address,
-        user,
-        password,
-        space,
-        vid_type,
-        scenarios,
-        vu,
-        duration,
-        dry_run,
-        **kwargs
+        self, folder, address, user, password, space, vid_type, scenarios, args, dry_run, **kwargs
     ):
         self.folder = folder or setting.DATA_FOLDER
         self.address = address or setting.NEBULA_ADDRESS
@@ -44,9 +36,8 @@ def __init__(
         self.vid_type = vid_type
         self.scenarios = []
         self.output_folder = "output"
-        self.vu = vu
-        self.duration = duration
         self.dry_run = dry_run
+        self.args = args
         self.scenarios = load_scenarios(scenarios)
         logger.info("total stress test scenarios is {}".format(len(self.scenarios)))
 
@@ -72,27 +63,18 @@ def gen_stress(
         space,
         vid_type,
         scenarios,
-        vu,
-        duration,
+        args,
         dry_run=None,
         **kwargs
     ):
         if _type.upper() not in cls.type_list:
             raise Exception("not impletment this test tool, tool is {}".format(_type))
 
         clazz = cls.get_all_stress_class().get("{}Stress".format(_type.upper()), None)
+        if args is not None:
+            args = args.strip()
         return clazz(
-            folder,
-            address,
-            user,
-            password,
-            space,
-            vid_type,
-            scenarios,
-            vu,
-            duration,
-            dry_run,
-            **kwargs
+            folder, address, user, password, space, vid_type, scenarios, args, dry_run, **kwargs
         )
 
     @classmethod
@@ -140,35 +122,99 @@ def dump_config(self, scenario):
         )
         jinja_dump(template_file, "{}/{}.js".format(self.output_folder, name), kwargs)
 
+    def _get_params(self):
+        """
+        e.g.
+        args:
+            "-s 60s:0 -s 40s:30 -v"
+        return:
+            {
+                "-s": ["60s:0", "40s:30"],
+                "-v": None
+            }
+        """
+        r = {}
+        if self.args is None:
+            return r
+
+        key, value = None, None
+        for item in self.args.split(" "):
+            if item.startswith("-"):
+                if key is not None and key not in r:
+                    r[key] = None
+                key = item
+            elif item.strip() != "":
+                value = item
+                if key not in r:
+                    r[key] = [value]
+                else:
+                    r[key].append(value)
+
+        if key is not None and key not in r:
+            r[key] = None
+        return r
+
     def run(self):
         logger.info("run stress test in k6")
-        logger.info(
-            "every scenario would run by {} vus and last {} secconds".format(self.vu, self.duration)
-        )
+        params = self._get_params()
+
+        # cannot use both stage and vu
+        run_with_stage = "-s" in params or "--stage" in params
+        vu = self.DEFAULT_VU
+        duration = self.DEFAULT_DURATION
+        if "-u" in params:
+            vu = params.pop("-u")[0]
+        if "--vus" in params:
+            vu = params.pop("--vus")[0]
+        if "-vu" in params:
+            vu = params.pop("-vu")[0]
+
+        if "-d" in params:
+            duration = params.pop("-d")[0]
+        if "--duration" in params:
+            duration = params.pop("--duration")[0]
+
+        logger.info("every scenario would run by {} vus and last {}".format(vu, duration))
         Path(self.output_folder).mkdir(exist_ok=True)
-        for scenario in self.scenarios:
+        if "--summary-trend-stats" not in params:
+            params["--summary-trend-stats"] = ["min,avg,med,max,p(90),p(95),p(99)"]
+        if setting.INFLUXDB_URL is not None and "--out" not in params and "-o" not in params:
+            params["--out"] = ["influxdb={}".format(setting.INFLUXDB_URL)]
 
+        for scenario in self.scenarios:
             self.dump_config(scenario)
-            # run k6
-            command = [
-                "scripts/k6",
-                "run",
-                "{}/{}.js".format(self.output_folder, scenario.name),
-                "-u",
-                str(self.vu),
-                "-d",
-                "{}s".format(self.duration),
-                "--summary-trend-stats",
-                "min,avg,med,max,p(90),p(95),p(99)",
-                "--summary-export",
-                "{}/result_{}.json".format(self.output_folder, scenario.name),
-            ]
-            if setting.INFLUXDB_URL is not None:
-                command.append("--out")
-                command.append("influxdb={}".format(setting.INFLUXDB_URL))
+            if run_with_stage:
+                command = [
+                    "scripts/k6",
+                    "run",
+                    "{}/{}.js".format(self.output_folder, scenario.name),
+                ]
+            else:
+                command = [
+                    "scripts/k6",
+                    "run",
+                    "{}/{}.js".format(self.output_folder, scenario.name),
+                    "-u",
+                    str(vu),
+                    "-d",
+                    str(duration),
+                ]
+
+            if "--summary-export" not in params:
+                params["--summary-export"] = [
+                    "{}/result_{}.json".format(self.output_folder, scenario.name)
+                ]
+
+            for param, values in params.items():
+                if values is None:
+                    command.append(param)
+                else:
+                    for v in values:
+                        command.append(param)
+                        command.append(v)
 
             click.echo("run command as below:")
-            click.echo(" ".join(command))
+            click.echo(" ".join([x if "(" not in x else '"{}"'.format(x) for x in command]))
             if self.dry_run is not None and self.dry_run:
                 continue
             run_process(command)