add ref_url

Author: erwin-wee

chainbench/main.py

@@ -210,6 +210,7 @@ def validate_profile_path(ctx: Context, param: Parameter, value: str) -> str:
 @click.option("--pg-password", default=None, help="PG password")
 @click.option("--use-latest-blocks", is_flag=True, help="Uses latest blocks for test data")
 @click.option("--size", default=None, help="Set the size of the test data. e.g. --size S")
+@click.option("--ref-url", default=None, help="Reference Node URL for retrieving test data before test starts. If not specified, target url is used instead.")
 @click.pass_context
 def start(
     ctx: Context,
@@ -243,6 +244,7 @@ def start(
     use_latest_blocks: bool,
     size: str | None,
     method: str | None = None,
+    ref_url: str | None = None,
 ) -> None:
     if notify:
         click.echo(f"Notify when test is finished using topic: {notify}")
@@ -280,7 +282,7 @@ def start(
 
     user_classes = {}
     for locustfile in parse_locustfile_paths([final_profile_path.__str__()]):
-        _, _user_classes, _ = load_locustfile(locustfile)
+        _user_classes, _  = load_locustfile(locustfile)
         for key, value in _user_classes.items():
             user_classes[key] = value
     test_data_types = set()
@@ -379,6 +381,7 @@ def start(
         method=method,
         enable_class_picker=enable_class_picker,
         batch_size=batch_size,
+        ref_url=ref_url,
     )
     # Start the Locust master
     master_command = locust_options.get_master_command()

chainbench/util/cli.py

@@ -144,6 +144,7 @@ class LocustOptions:
     method: str | None = None
     enable_class_picker: bool = False
     batch_size: int | None = None
+    ref_url: str | None = None
 
     def get_master_command(self) -> str:
         """Generate master command."""
@@ -203,6 +204,9 @@ def get_extra_options(self, command: str):
 
         if self.batch_size is not None:
             command += f" --batch-size {self.batch_size}"
+
+        if self.ref_url is not None:
+            command += f" --ref-url {self.ref_url}"
         return command
 
 

chainbench/util/event.py

@@ -31,7 +31,7 @@ def cli_custom_arguments(parser: LocustArgumentParser):
         "--size",
         type=str,
         default=None,
-        help="Set the size of the test data. e.g. --size S",
+        help="Set the size of the test data. e.g. --size S.",
         include_in_web_ui=False,
     )
     parser.add_argument(
@@ -48,7 +48,14 @@ def cli_custom_arguments(parser: LocustArgumentParser):
         type=str,
         default="eth_blockNumber",
         choices=list(all_methods.keys()),
-        help="Test a specific method",
+        help="Test a specific method.",
+        include_in_web_ui=True,
+    )
+    parser.add_argument(
+        "--ref-url",
+        type=str,
+        default=None,
+        help="Set the reference node url to be used for retrieving test data before test starts. If empty, defaults to target url.",
         include_in_web_ui=True,
     )
 
@@ -187,15 +194,15 @@ def on_init(environment: Environment, **_kwargs):
                         continue
                     logger.info(f"Initializing test data for {test_data_class_name}")
                     print(f"Initializing test data for {test_data_class_name}")
-                    if environment.host:
-                        user_test_data.init_http_client(environment.host)
-                    if isinstance(user_test_data, EvmTestData):
-                        chain_id: ChainId = user_test_data.fetch_chain_id()
-                        user_test_data.init_network(chain_id)
-                        logger.info(f"Target endpoint network is {user_test_data.network.name}")
-                        print(f"Target endpoint network is {user_test_data.network.name}")
-                        test_data["chain_id"] = {test_data_class_name: chain_id}
                     if environment.parsed_options:
+                        ref_node_url = environment.parsed_options.ref_url if environment.parsed_options.ref_url is not None else environment.host
+                        user_test_data.init_http_client(ref_node_url)
+                        if isinstance(user_test_data, EvmTestData):
+                            chain_id: ChainId = user_test_data.fetch_chain_id()
+                            user_test_data.init_network(chain_id)
+                            logger.info(f"Target endpoint network is {user_test_data.network.name}")
+                            print(f"Target endpoint network is {user_test_data.network.name}")
+                            test_data["chain_id"] = {test_data_class_name: chain_id}
                         user_test_data.init_data(environment.parsed_options)
                         test_data[test_data_class_name] = user_test_data.data.to_json()
                         send_msg_to_workers(environment.runner, "test_data", test_data)