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

@@ -69,7 +69,13 @@ const EXAMPLE_L1_HANDLER_TX_HASH: &str =
 /// Retrieves the test URL from the `TEST_URL` environment variable,
 /// falling back to a default URL if not provided.
 fn get_test_url() -> String {
-    RPC_NODE_URL.clone()
+    let url = RPC_NODE_URL.clone();
+    assert_eq!(
+        guess_chain_id_from_node_url(&url).unwrap(),
+        ChainId::Mainnet,
+        "RPC HTTP tests not supported on chains other than mainnet."
+    );
+    url
 }
 
 /// Retrieves the test block_number from the `TEST_URL` environment variable,
@@ -88,24 +94,14 @@ pub fn get_test_rpc_config() -> RpcStateReaderConfig {
 }
 
 #[fixture]
-pub fn test_url() -> String {
-    get_test_url()
-}
-
-#[fixture]
-pub fn test_chain_id(test_url: String) -> ChainId {
-    guess_chain_id_from_node_url(test_url.as_str()).unwrap()
-}
-
-#[fixture]
-pub fn test_state_reader(test_chain_id: ChainId) -> TestStateReader {
+pub fn test_state_reader() -> TestStateReader {
     TestStateReader {
         rpc_state_reader: RpcStateReader {
             config: get_test_rpc_config(),
             block_id: get_test_block_id(),
         },
         retry_config: RetryConfig::default(),
-        chain_id: test_chain_id,
+        chain_id: ChainId::Mainnet,
         contract_class_mapping_dumper: Arc::new(Mutex::new(None)),
     }
 }
@@ -123,9 +119,8 @@ pub fn last_constructed_block(test_block_number: BlockNumber) -> BlockNumber {
 #[fixture]
 pub fn test_state_readers_last_and_current_block(
     last_constructed_block: BlockNumber,
-    test_chain_id: ChainId,
 ) -> ConsecutiveTestStateReaders {
-    ConsecutiveTestStateReaders::new(last_constructed_block, None, test_chain_id, false)
+    ConsecutiveTestStateReaders::new(last_constructed_block, None, ChainId::Mainnet, false)
 }
 
 /// Test that the block info can be retrieved from the RPC server.
@@ -244,7 +239,7 @@ pub fn test_get_statediff_rpc(test_state_reader: TestStateReader) {
 }
 
 #[rstest]
-#[case(test_block_number(test_state_reader(test_chain_id(test_url()))).0)]
+#[case(test_block_number(test_state_reader()).0)]
 #[case(EXAMPLE_DECLARE_V1_BLOCK_NUMBER)]
 #[case(EXAMPLE_DECLARE_V2_BLOCK_NUMBER)]
 #[case(EXAMPLE_DECLARE_V3_BLOCK_NUMBER)]