From 716bed2510ef308493deb8ace955549756c13965 Mon Sep 17 00:00:00 2001
From: Dori Medini <dori@starkware.co>
Date: Tue, 12 Nov 2024 09:50:08 +0200
Subject: [PATCH] refactor(blockifier_reexecution): inject chain ID to chain
 info creation

---
 .../src/state_reader/test_state_reader.rs        |  6 +++---
 .../src/state_reader/utils.rs                    | 16 ++++++++++------
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/crates/blockifier_reexecution/src/state_reader/test_state_reader.rs b/crates/blockifier_reexecution/src/state_reader/test_state_reader.rs
index c26e20d5b0..ae697b01c6 100644
--- a/crates/blockifier_reexecution/src/state_reader/test_state_reader.rs
+++ b/crates/blockifier_reexecution/src/state_reader/test_state_reader.rs
@@ -17,7 +17,7 @@ use blockifier::versioned_constants::VersionedConstants;
 use serde::{Deserialize, Serialize};
 use serde_json::{json, to_value};
 use starknet_api::block::{BlockHash, BlockHashAndNumber, BlockNumber, StarknetVersion};
-use starknet_api::core::{ClassHash, CompiledClassHash, ContractAddress, Nonce};
+use starknet_api::core::{ChainId, ClassHash, CompiledClassHash, ContractAddress, Nonce};
 use starknet_api::state::StorageKey;
 use starknet_api::transaction::{Transaction, TransactionHash};
 use starknet_core::types::ContractClass as StarknetContractClass;
@@ -131,7 +131,7 @@ impl From<SerializableOfflineReexecutionData> for OfflineReexecutionData {
             offline_state_reader_prev_block,
             block_context_next_block: BlockContext::new(
                 block_info_next_block,
-                get_chain_info(),
+                get_chain_info(&ChainId::Mainnet),
                 VersionedConstants::get(&starknet_version).unwrap().clone(),
                 BouncerConfig::max(),
             ),
@@ -327,7 +327,7 @@ impl TestStateReader {
     pub fn get_block_context(&self) -> ReexecutionResult<BlockContext> {
         Ok(BlockContext::new(
             self.get_block_info()?,
-            get_chain_info(),
+            get_chain_info(&ChainId::Mainnet),
             self.get_versioned_constants()?.clone(),
             BouncerConfig::max(),
         ))
diff --git a/crates/blockifier_reexecution/src/state_reader/utils.rs b/crates/blockifier_reexecution/src/state_reader/utils.rs
index b5a8aab831..f10ba3afe5 100644
--- a/crates/blockifier_reexecution/src/state_reader/utils.rs
+++ b/crates/blockifier_reexecution/src/state_reader/utils.rs
@@ -25,10 +25,14 @@ pub const RPC_NODE_URL: &str = "https://free-rpc.nethermind.io/mainnet-juno/";
 pub const JSON_RPC_VERSION: &str = "2.0";
 
 /// Returns the fee token addresses of mainnet.
-pub fn get_fee_token_addresses() -> FeeTokenAddresses {
-    FeeTokenAddresses {
-        strk_fee_token_address: *STRK_FEE_CONTRACT_ADDRESS,
-        eth_fee_token_address: *ETH_FEE_CONTRACT_ADDRESS,
+pub fn get_fee_token_addresses(chain_id: &ChainId) -> FeeTokenAddresses {
+    match chain_id {
+        // Mainnet, testnet and integration systems have the same fee token addresses.
+        ChainId::Mainnet | ChainId::Sepolia | ChainId::IntegrationSepolia => FeeTokenAddresses {
+            strk_fee_token_address: *STRK_FEE_CONTRACT_ADDRESS,
+            eth_fee_token_address: *ETH_FEE_CONTRACT_ADDRESS,
+        },
+        unknown_chain => unimplemented!("Unknown chain ID {unknown_chain}."),
     }
 }
 
@@ -41,8 +45,8 @@ pub fn get_rpc_state_reader_config() -> RpcStateReaderConfig {
 }
 
 /// Returns the chain info of mainnet.
-pub fn get_chain_info() -> ChainInfo {
-    ChainInfo { chain_id: ChainId::Mainnet, fee_token_addresses: get_fee_token_addresses() }
+pub fn get_chain_info(chain_id: &ChainId) -> ChainInfo {
+    ChainInfo { chain_id: chain_id.clone(), fee_token_addresses: get_fee_token_addresses(chain_id) }
 }
 
 // TODO(Aner): import the following functions instead, to reduce code duplication.