add chain parameter

galacticcouncil · Mar 22, 2022 · 3d1c2ce · 3d1c2ce
1 parent 23c5fbf
commit 3d1c2ce
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/bench_wizard/benchmark.py b/bench_wizard/benchmark.py
@@ -16,6 +16,7 @@ class BenchmarksConfig:
     dump_results: Optional[str] = None
     output_dir: Optional[str] = None
     template: Optional[str] = None
+    chain: Optional[str] = "dev"
 
 
 class Benchmark:
@@ -77,7 +78,7 @@ def _prepare_benchmarks(config: BenchmarksConfig) -> List[Benchmark]:
     benchmarks = []
 
     for pallet in config.pallets:
-        cargo = Cargo(pallet=pallet, template=config.template)
+        cargo = Cargo(pallet=pallet, template=config.template, chain=config.chain)
 
         if config.output_dir:
             output_file = os.path.join(config.output_dir, f"{pallet}.rs")

diff --git a/bench_wizard/main.py b/bench_wizard/main.py
@@ -31,6 +31,14 @@ def version():
     required=True,
     help="Pallets",
 )
+@click.option(
+    "-c",
+    "--chain",
+    type=str,
+    required=False,
+    default="dev",
+    help="chain",
+)
 @click.option(
     "-d",
     "--dump-results",
@@ -54,6 +62,7 @@ def version():
 )
 def benchmark(
     pallet: list,
+    chain: str,
     dump_results: Optional[str],
     template: Optional[str],
     output_dir: Optional[str],
@@ -64,6 +73,7 @@ def benchmark(
         dump_results=dump_results,
         template=template,
         output_dir=output_dir,
+        chain=chain,
     )
 
     run_pallet_benchmarks(config, Output())
@@ -94,7 +104,6 @@ def pc(
         print(f"{reference_values} does not exist", file=sys.stderr)
         exit(1)
 
-
     config = PerformanceConfig(
         reference_values=reference_values,
         pallets=pallet,