From 3d1c2ce3fd80169b84a0a06599575eaeb1f731f9 Mon Sep 17 00:00:00 2001
From: Martin <martin.hloska@gmail.com>
Date: Tue, 22 Mar 2022 15:04:49 +0100
Subject: [PATCH] add chain parameter

---
 bench_wizard/benchmark.py |  3 ++-
 bench_wizard/main.py      | 11 ++++++++++-
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/bench_wizard/benchmark.py b/bench_wizard/benchmark.py
index 841cf2c..8f580bc 100644
--- 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
index 458f670..362e024 100644
--- 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,