refactoring

- fix mypy errors
- added eth_estimate_gas_task to EVMMethods
- deprecated _eth_estimate_gas_params_factory

Author: erwin-wee

chainbench/profile/evm/heavy.py

@@ -58,7 +58,7 @@ def trace_replay_block_transactions_task(self):
         self.make_call(
             name="trace_replay_block_transactions",
             method="trace_replayBlockTransactions",
-            params=self._trace_replay_block_transaction_params_factory(get_rng()),
+            params=self._trace_replay_block_transaction_params_factory(),
         )
 
     @task
@@ -137,5 +137,5 @@ def eth_estimate_gas_task(self):
         self.make_call(
             name="eth_estimate_gas",
             method="eth_estimateGas",
-            params=self._eth_estimate_gas_params_factory(get_rng()),
+            params=self._erc20_eth_call_params_factory(get_rng()),
         )

chainbench/test_data/blockchain/evm.py

@@ -6,35 +6,23 @@
 
 
 class SmartContract:
-    def __init__(self, address: str, function_params: list[t.Callable]):
+    def __init__(self, address: str):
         self.address = address
-        self.function_params = function_params
 
 
 class ERC20Contract(SmartContract):
-    def __init__(self, address: str):
-        function_params = [
-            self.total_supply_params,
-            self.symbol_params,
-            self.balance_of_params,
-            self.name_params,
-        ]
-        super().__init__(address, function_params)
-
-    def total_supply_params(self, **kwargs) -> dict[str, str]:
+    def total_supply_params(self) -> dict[str, str]:
         return {"data": "0x18160ddd", "to": self.address}
 
-    def balance_of_params(self, **kwargs) -> dict[str, str]:
-        return {"data": "0x70a08231" + kwargs["address"][2:].zfill(64), "to": self.address}
+    def balance_of_params(self, target_address: str) -> dict[str, str]:
+        return {"data": "0x70a08231" + target_address[2:].zfill(64), "to": self.address}
 
-    def symbol_params(self, **kwargs) -> dict[str, str]:
+    def symbol_params(self) -> dict[str, str]:
         return {"data": "0x95d89b41", "to": self.address}
 
-    def name_params(self, **kwargs) -> dict[str, str]:
+    def name_params(self) -> dict[str, str]:
         return {"data": "0x06fdde03", "to": self.address}
 
-    def get_random_function_params(self, rng: RNG) -> t.Callable:
-        return rng.random.choice(self.function_params)
 
 class NetworkInfo(t.TypedDict):
     name: str

chainbench/test_data/evm.py

@@ -13,7 +13,8 @@
     Tx,
     TxHash,
 )
-from chainbench.test_data.blockchain.evm import ChainInfo
+from chainbench.test_data.blockchain.evm import ChainInfo, ERC20Contract
+from chainbench.util.rng import RNG
 
 logger = logging.getLogger(__name__)
 
@@ -26,6 +27,12 @@ def __init__(self, rpc_version: str = "2.0"):
     def set_chain_info(self, chain_id: ChainId):
         self.chain_info = ChainInfo(chain_id)
 
+    def get_random_erc20_contract(self, rng: RNG) -> ERC20Contract:
+        if self.chain_info is None:
+            raise ValueError("Chain info is not set")
+        else:
+            return self.chain_info.get_random_contract(rng)
+
     def fetch_chain_id(self) -> int:
         return self._parse_hex_to_int(self._make_call("eth_chainId"))
 

chainbench/user/evm.py

@@ -79,7 +79,8 @@ def _trace_block_by_hash_params_factory(self, rng: RNG) -> list[BlockHash | dict
             {"tracer": "callTracer", "timeout": self._default_trace_timeout},
         ]
 
-    def _trace_replay_block_transaction_params_factory(self, rng: RNG) -> list[str | list[str]]:
+    @staticmethod
+    def _trace_replay_block_transaction_params_factory() -> list[str | list[str]]:
         return [
             "latest",
             ["vmTrace", "trace", "stateDiff"],
@@ -107,31 +108,30 @@ def _trace_filter_params_factory(self, rng: RNG) -> list[dict]:
             }
         ]
 
-    def _eth_estimate_gas_params_factory(self, rng: RNG) -> list[dict]:
-        return [
-            {
-                "from": self.test_data.get_random_account(rng),
-                "to": self.test_data.get_random_account(rng),
-            }
-        ]
-
-    def _eth_fee_history_params_factory(self, rng: RNG) -> list[int | str | list[int]]:
+    @staticmethod
+    def _eth_fee_history_params_factory(rng: RNG) -> list[int | str | list[int]]:
         return [rng.random.randint(1, 1024), "latest", [25, 75]]
 
     def _erc20_eth_call_params_factory(self, rng: RNG):
-        contract = self.test_data.chain_info.get_random_contract(rng)
-        function_params = contract.get_random_function_params(rng)
+        contract = self.test_data.get_random_erc20_contract(rng)
+        functions_params = [
+            contract.total_supply_params(),
+            contract.balance_of_params(self.test_data.get_random_account(rng)),
+            contract.symbol_params(),
+            contract.name_params(),
+        ]
         return [
-            function_params(address=self.test_data.get_random_account(rng)),
+            rng.random.choice(functions_params),
             "latest",
         ]
 
-    def _erc20_eth_get_code_params_factory(self, rng:RNG):
+    def _erc20_eth_get_code_params_factory(self, rng: RNG):
         return [
-            self.test_data.chain_info.get_random_contract(rng).address,
+            self.test_data.get_random_erc20_contract(rng).address,
             "latest",
         ]
 
+
 class EVMMethods(EVMBenchUser):
     abstract = True
 
@@ -156,6 +156,13 @@ def eth_chain_id_task(self):
             method="eth_chainId",
         )
 
+    def eth_estimate_gas_task(self):
+        self.make_call(
+            name="eth_estimate_gas",
+            method="eth_estimateGas",
+            params=self._erc20_eth_call_params_factory(self.rng.get_rng()),
+        )
+
     def eth_fee_history_task(self) -> None:
         self.make_call(
             method="eth_feeHistory",
@@ -333,7 +340,7 @@ def trace_block_task(self) -> None:
     def trace_replay_block_transactions_task(self) -> None:
         self.make_call(
             method="trace_replayBlockTransactions",
-            params=self._trace_replay_block_transaction_params_factory(self.rng.get_rng()),
+            params=self._trace_replay_block_transaction_params_factory(),
         )
 
     def trace_replay_transaction_task(self) -> None: