add --use-recent-blocks flag (#33)

* add --use-recent-blocks flag

for testing nodes that return null for certain historical blocks, --use-recent-blocks flag is added to use recent blocks as test data

* Update README.md

* fix mypy issues

* optimize test_data initialization

Author: erwin-wee

README.md

@@ -104,7 +104,7 @@ After the test is finished, the tool will automatically quit.
 - `--help`: Displays the help message.
 - `--debug-trace-methods`: Enables tasks tagged with debug or trace to be executed
 - `-E, --exclude-tags`: Exclude tasks tagged with custom tags from the test. You may specify this option multiple times --help Show this message and exit.
-
+- `--use-recent-blocks`: Use recent blocks for test data generation.
 You may also run `chainbench start --help` for the full list of parameters and flags.
 
 ### Profiles

chainbench/main.py

@@ -136,6 +136,9 @@ def cli(ctx: click.Context):
     "--pg-username", default="postgres", help="PG username", show_default=True
 )
 @click.option("--pg-password", default=None, help="PG password")
+@click.option(
+    "--use-recent-blocks", is_flag=True, help="Uses recent blocks for test data"
+)
 @click.pass_context
 def start(
     ctx: click.Context,
@@ -162,6 +165,7 @@ def start(
     pg_port: int,
     pg_username: str,
     pg_password: str | None,
+    use_recent_blocks: bool,
 ):
     if notify:
         click.echo(f"Notify when test is finished using topic: {notify}")
@@ -230,6 +234,7 @@ def start(
         pg_port=pg_port,
         pg_username=pg_username,
         pg_password=pg_password,
+        use_recent_blocks=use_recent_blocks,
     )
     if headless:
         click.echo(f"Starting master in headless mode for {profile}")
@@ -259,6 +264,7 @@ def start(
             pg_port=pg_port,
             pg_username=pg_username,
             pg_password=pg_password,
+            use_recent_blocks=use_recent_blocks,
         )
         worker_args = shlex.split(worker_command, posix=is_posix)
         worker_process = subprocess.Popen(worker_args)

chainbench/test_data/base.py

@@ -44,11 +44,11 @@ def __init__(self, rpc_version: str = "2.0"):
 
         self._data: BlockchainData | None = None
 
-    def update(self, host_url: str):
+    def update(self, host_url: str, use_recent_blocks: bool = False) -> BlockchainData:
         self._logger.info("Updating data")
         self._host = host_url
         self._logger.debug("Host: %s", self._host)
-        data = self._get_init_data()
+        data = self._get_init_data(use_recent_blocks)
         self._logger.info("Data fetched")
         self._logger.debug("Data: %s", data)
         self._data = data
@@ -57,7 +57,7 @@ def update(self, host_url: str):
         self._logger.info("Lock released")
         return data
 
-    def _get_init_data(self) -> BlockchainData:
+    def _get_init_data(self, use_recent_blocks) -> BlockchainData:
         raise NotImplementedError
 
     @property

chainbench/test_data/evm.py

@@ -1,6 +1,6 @@
 from typing import Mapping, TypedDict
 
-from chainbench.test_data.base import BaseTestData, BlockchainData
+from chainbench.test_data.base import BaseTestData, BlockchainData, Blocks
 from chainbench.util.rng import get_rng
 
 
@@ -11,9 +11,9 @@ class ChainInfo(TypedDict):
 
 
 class EVMTestData(BaseTestData):
-    TXS_REQUIRED = 50
-    ACCOUNTS_REQUIRED = 100
-    SAVE_LAST_BLOCKS = 100
+    TXS_REQUIRED = 100
+    ACCOUNTS_REQUIRED = 200
+    SAVE_BLOCKS = 20
 
     CHAIN_INFO: Mapping[int, ChainInfo] = {
         1: {
@@ -78,25 +78,45 @@ def _fetch_random_block(self, start, end, return_txs=True) -> tuple[int, dict]:
         return self._fetch_block(block_number, return_txs=return_txs)
 
     # get initial data from blockchain
-    def _get_init_data(self) -> BlockchainData:
+    def _get_init_data(self, use_recent_blocks) -> BlockchainData:
         txs: list[dict] = []
         tx_hashes: list[str] = []
         accounts: set[str] = set()
-        blocks = []
-        chain_id = self._fetch_chain_id()
-        start_block_number = self.CHAIN_INFO[chain_id]["start_block"]
-        end_block_number = self.CHAIN_INFO[chain_id]["end_block"]
+        blocks: Blocks = []
+        chain_id: int = self._fetch_chain_id()
+        start_block_number: int
+        end_block_number: int
+        return_txs: bool
 
-        while self.TXS_REQUIRED > len(txs) and self.ACCOUNTS_REQUIRED > len(accounts):
+        if use_recent_blocks:
+            end_block_number = int(self._make_call("eth_blockNumber"), 0)
+            start_block_number = end_block_number - 20
+        else:
+            start_block_number = self.CHAIN_INFO[chain_id]["start_block"]
+            end_block_number = self.CHAIN_INFO[chain_id]["end_block"]
+
+        while (
+            self.TXS_REQUIRED > len(txs)
+            or self.ACCOUNTS_REQUIRED > len(accounts)
+            or self.SAVE_BLOCKS > len(blocks)
+        ):
+            if self.ACCOUNTS_REQUIRED > len(accounts) or self.SAVE_BLOCKS > len(blocks):
+                return_txs = True
+            else:
+                return_txs = False
             block_number, block = self._fetch_random_block(
-                start_block_number, end_block_number
+                start_block_number, end_block_number, return_txs
             )
-            blocks.append((block_number, block["hash"]))
+            if self.SAVE_BLOCKS > len(blocks):
+                blocks.append((block_number, block["hash"]))
             for tx in block["transactions"]:
-                self._append_if_not_none(txs, tx)
-                self._append_if_not_none(tx_hashes, tx["hash"])
-                self._append_if_not_none(accounts, tx["from"])
-                self._append_if_not_none(accounts, tx["to"])
+                if self.TXS_REQUIRED > len(txs):
+                    self._append_if_not_none(txs, tx)
+                    self._append_if_not_none(tx_hashes, tx["hash"])
+                if self.ACCOUNTS_REQUIRED > len(accounts):
+                    self._append_if_not_none(accounts, tx["from"])
+                if self.ACCOUNTS_REQUIRED > len(accounts):
+                    self._append_if_not_none(accounts, tx["to"])
 
         return BlockchainData(
             end_block_number=end_block_number,

chainbench/util/cli.py

@@ -70,6 +70,7 @@ def get_master_command(
     pg_port: int | None = None,
     pg_username: str | None = None,
     pg_password: str | None = None,
+    use_recent_blocks: bool = False,
 ) -> str:
     """Generate master command."""
     command = (
@@ -94,6 +95,8 @@ def get_master_command(
     if len(exclude_tags) > 0:
         command += f" --exclude-tags {' '.join(exclude_tags)}"
 
+    if use_recent_blocks:
+        command += " --use-recent-blocks True"
     return command
 
 
@@ -112,6 +115,7 @@ def get_worker_command(
     pg_port: int | None = None,
     pg_username: str | None = None,
     pg_password: str | None = None,
+    use_recent_blocks: bool = False,
 ) -> str:
     """Generate worker command."""
     command = (
@@ -131,6 +135,8 @@ def get_worker_command(
     if len(exclude_tags) > 0:
         command += f" --exclude-tags {' '.join(exclude_tags)}"
 
+    if use_recent_blocks:
+        command += " --use-recent-blocks True"
     return command
 
 

chainbench/util/event.py

@@ -8,6 +8,42 @@
 logger = logging.getLogger(__name__)
 
 
+def cli_custom_arguments(parser):
+    parser.add_argument(
+        "--use-recent-blocks",
+        type=bool,
+        default=False,
+        help="Use recent blocks as test data",
+    )
+
+
+# Listener for the init event
+def on_init(environment, **_kwargs):
+    # It will be called for any runner (master, worker, local)
+    logger.debug("init.add_listener: Init is started")
+    logger.debug("init.add_listener: Environment: %s", environment.runner)
+    logger.debug("init.add_listener: Host: %s", environment.host)
+
+    host_under_test = environment.host or "Default host"
+
+    if isinstance(environment.runner, MasterRunner):
+        # Print master details to the log
+        logger.info("I'm a master. Running tests for %s", host_under_test)
+
+    if isinstance(environment.runner, WorkerRunner):
+        # Print worker details to the log
+        logger.info("I'm a worker. Running tests for %s", host_under_test)
+        logger.info("Initializing test data...")
+        for user in environment.runner.user_classes:
+            if not hasattr(user, "test_data"):
+                continue
+
+            user.test_data.update(
+                environment.host, environment.parsed_options.use_recent_blocks
+            )
+            logger.info("Test data is ready")
+
+
 def on_test_start(environment, **_kwargs):
     # It will be called for any runner (master, worker, local)
     if not isinstance(environment.runner, MasterRunner):
@@ -40,32 +76,8 @@ def on_test_stop(environment, **_kwargs):
         logger.info("Master: The test is stopped")
 
 
-# Listener for the init event
-def on_init(environment, **_kwargs):
-    # It will be called for any runner (master, worker, local)
-    logger.debug("init.add_listener: Init is started")
-    logger.debug("init.add_listener: Environment: %s", environment.runner)
-    logger.debug("init.add_listener: Host: %s", environment.host)
-
-    host_under_test = environment.host or "Default host"
-
-    if isinstance(environment.runner, MasterRunner):
-        # Print master details to the log
-        logger.info("I'm a master. Running tests for %s", host_under_test)
-
-    if isinstance(environment.runner, WorkerRunner):
-        # Print worker details to the log
-        logger.info("I'm a worker. Running tests for %s", host_under_test)
-        logger.info("Initializing test data...")
-        for user in environment.runner.user_classes:
-            if not hasattr(user, "test_data"):
-                continue
-
-            user.test_data.update(environment.host)
-            logger.info("Test data is ready")
-
-
 def setup_event_listeners():
+    events.init_command_line_parser.add_listener(cli_custom_arguments)
     events.test_start.add_listener(on_test_start)
     events.test_stop.add_listener(on_test_stop)
     events.init.add_listener(on_init)