refactor(blockifier_reexecution): allow overriding chain ID explicitly

crates/blockifier_reexecution/src/main.rs

@@ -11,6 +11,7 @@ use blockifier_reexecution::state_reader::utils::{
 };
 use clap::{Args, Parser, Subcommand};
 use starknet_api::block::BlockNumber;
+use starknet_api::core::ChainId;
 use starknet_gateway::config::RpcStateReaderConfig;
 
 /// BlockifierReexecution CLI.
@@ -24,13 +25,49 @@ pub struct BlockifierReexecutionCliArgs {
     command: Command,
 }
 
+#[derive(clap::ValueEnum, Clone, Debug)]
+enum SupportedChainId {
+    Mainnet,
+    Testnet,
+    Integration,
+}
+
+impl From<SupportedChainId> for ChainId {
+    fn from(chain_id: SupportedChainId) -> Self {
+        match chain_id {
+            SupportedChainId::Mainnet => Self::Mainnet,
+            SupportedChainId::Testnet => Self::Sepolia,
+            SupportedChainId::Integration => Self::IntegrationSepolia,
+        }
+    }
+}
+
+#[derive(Debug, Args)]
+struct RpcArgs {
+    /// Node url.
+    #[clap(long, short = 'n')]
+    node_url: String,
+
+    /// Optional chain ID (if not provided, it will be guessed from the node url).
+    #[clap(long, short = 'c')]
+    chain_id: Option<SupportedChainId>,
+}
+
+impl RpcArgs {
+    pub(crate) fn parse_chain_id(&self) -> ChainId {
+        self.chain_id
+            .clone()
+            .map(ChainId::from)
+            .unwrap_or(guess_chain_id_from_node_url(self.node_url.as_str()).unwrap())
+    }
+}
+
 #[derive(Debug, Subcommand)]
 enum Command {
     /// Runs the RPC test.
     RpcTest {
-        /// Node url.
-        #[clap(long, short = 'n')]
-        node_url: String,
+        #[clap(flatten)]
+        rpc_args: RpcArgs,
 
         /// Block number.
         #[clap(long, short = 'b')]
@@ -39,9 +76,8 @@ enum Command {
 
     /// Writes the RPC queries to json files.
     WriteRpcRepliesToJson {
-        /// Node url.
-        #[clap(long, short = 'n')]
-        node_url: String,
+        #[clap(flatten)]
+        rpc_args: RpcArgs,
 
         /// Block number.
         #[clap(long, short = 'b')]
@@ -56,9 +92,8 @@ enum Command {
 
     /// Writes the RPC queries of all (selected) blocks to json files.
     WriteAll {
-        /// Node url.
-        #[clap(long, short = 'n')]
-        node_url: String,
+        #[clap(flatten)]
+        rpc_args: RpcArgs,
 
         /// Block numbers. If not specified, blocks are retrieved from
         /// get_block_numbers_for_reexecution().
@@ -95,18 +130,21 @@ fn main() {
     let args = BlockifierReexecutionCliArgs::parse();
 
     match args.command {
-        Command::RpcTest { node_url, block_number } => {
-            println!("Running RPC test for block number {block_number} using node url {node_url}.",);
+        Command::RpcTest { block_number, rpc_args } => {
+            println!(
+                "Running RPC test for block number {block_number} using node url {}.",
+                rpc_args.node_url
+            );
 
             let config = RpcStateReaderConfig {
-                url: node_url.clone(),
+                url: rpc_args.node_url.clone(),
                 json_rpc_version: JSON_RPC_VERSION.to_string(),
             };
 
             reexecute_and_verify_correctness(ConsecutiveTestStateReaders::new(
                 BlockNumber(block_number - 1),
                 Some(config),
-                guess_chain_id_from_node_url(node_url.as_str()).unwrap(),
+                rpc_args.parse_chain_id(),
                 false,
             ));
 
@@ -115,7 +153,7 @@ fn main() {
             println!("RPC test passed successfully.");
         }
 
-        Command::WriteRpcRepliesToJson { node_url, block_number, full_file_path } => {
+        Command::WriteRpcRepliesToJson { block_number, full_file_path, rpc_args } => {
             let full_file_path = full_file_path.unwrap_or(format!(
                 "./crates/blockifier_reexecution/resources/block_{block_number}/reexecution_data.\
                  json"
@@ -124,12 +162,12 @@ fn main() {
             write_block_reexecution_data_to_file(
                 BlockNumber(block_number),
                 &full_file_path,
-                node_url.clone(),
-                guess_chain_id_from_node_url(node_url.as_str()).unwrap(),
+                rpc_args.node_url.clone(),
+                rpc_args.parse_chain_id(),
             );
         }
 
-        Command::WriteAll { node_url, block_numbers, directory_path } => {
+        Command::WriteAll { block_numbers, directory_path, rpc_args } => {
             let directory_path =
                 directory_path.unwrap_or("./crates/blockifier_reexecution/resources".to_string());
 
@@ -148,8 +186,8 @@ fn main() {
                 write_block_reexecution_data_to_file(
                     block_number,
                     &full_file_path,
-                    node_url.clone(),
-                    guess_chain_id_from_node_url(node_url.as_str()).unwrap(),
+                    rpc_args.node_url.clone(),
+                    rpc_args.parse_chain_id(),
                 );
             }
         }

crates/blockifier_reexecution/src/state_reader/rpc_https_test.rs

@@ -36,8 +36,12 @@ use crate::state_reader::reexecution_state_reader::ReexecutionStateReader;
 use crate::state_reader::test_state_reader::{ConsecutiveTestStateReaders, TestStateReader};
 use crate::state_reader::utils::guess_chain_id_from_node_url;
 
-const EXAMPLE_INVOKE_TX_HASH: &str =
+const EXAMPLE_INVOKE_TX_HASH_MAINNET: &str =
     "0xa7c7db686c7f756ceb7ca85a759caef879d425d156da83d6a836f86851983";
+const EXAMPLE_INVOKE_TX_HASH_TESTNET: &str =
+    "0x008dcac0b86876435eb88fad26f45b8251009db008061ff15154de3f504a7e7c";
+const EXAMPLE_INVOKE_TX_HASH_INTEGRATION: &str =
+    "0x4b1433f4c92710b6bb3a2b9304ec5169141227a1c25ddd2611938fc52964ba5";
 
 const EXAMPLE_CONTACT_CLASS_HASH: &str =
     "0x3131fa018d520a037686ce3efddeab8f28895662f019ca3ca18a626650f7d1e";
@@ -97,6 +101,16 @@ pub fn test_chain_id(test_url: String) -> ChainId {
     guess_chain_id_from_node_url(test_url.as_str()).unwrap()
 }
 
+#[fixture]
+pub fn example_invoke_tx_hash(test_chain_id: ChainId) -> &'static str {
+    match test_chain_id {
+        ChainId::Mainnet => EXAMPLE_INVOKE_TX_HASH_MAINNET,
+        ChainId::Sepolia => EXAMPLE_INVOKE_TX_HASH_TESTNET,
+        ChainId::IntegrationSepolia => EXAMPLE_INVOKE_TX_HASH_INTEGRATION,
+        ChainId::Other(other) => panic!("Unknown chain id: {other}"),
+    }
+}
+
 #[fixture]
 pub fn test_state_reader(test_chain_id: ChainId) -> TestStateReader {
     TestStateReader {
@@ -173,8 +187,11 @@ pub fn test_get_tx_hashes(test_state_reader: TestStateReader) {
 }
 
 #[rstest]
-pub fn test_get_invoke_tx_by_hash(test_state_reader: TestStateReader) {
-    let actual_tx = test_state_reader.get_tx_by_hash(EXAMPLE_INVOKE_TX_HASH).unwrap();
+pub fn test_get_invoke_tx_by_hash(
+    test_state_reader: TestStateReader,
+    example_invoke_tx_hash: &str,
+) {
+    let actual_tx = test_state_reader.get_tx_by_hash(example_invoke_tx_hash).unwrap();
     assert_matches!(actual_tx, Transaction::Invoke(..));
 }