From 5a45aa7c1a0031eb7bb1601b0930da15991e9606 Mon Sep 17 00:00:00 2001
From: varun-doshi <doshivarun202@gmail.com>
Date: Thu, 26 Dec 2024 16:59:54 +0530
Subject: [PATCH] feat:remove duplication

---
 crates/pevm/tests/mainnet.rs | 85 +++++++++++++++++++++++-------------
 1 file changed, 54 insertions(+), 31 deletions(-)

diff --git a/crates/pevm/tests/mainnet.rs b/crates/pevm/tests/mainnet.rs
index 61fdd14d..37bc2aa4 100644
--- a/crates/pevm/tests/mainnet.rs
+++ b/crates/pevm/tests/mainnet.rs
@@ -1,15 +1,59 @@
 //! Test with mainnet blocks
-
-use pevm::chain::PevmEthereum;
+use reqwest::Url;
 
 pub mod common;
 
-#[tokio::test(flavor = "multi_thread")]
+enum Chain {
+    Eth,
+    Op,
+}
 #[cfg(feature = "rpc-storage")]
-async fn mainnet_blocks_from_rpc() {
+async fn blocks_from_rpc(rpc_url: &Url, block_number: u64, chain_name: Chain) {
     use alloy_provider::{Provider, ProviderBuilder};
     use alloy_rpc_types_eth::{BlockId, BlockTransactionsKind};
     use pevm::chain::PevmChain;
+    
+    
+    match chain_name {
+        Chain::Eth => {
+            use pevm::chain::PevmEthereum;
+            let provider = ProviderBuilder::new().on_http(rpc_url.clone());
+            let chain = PevmEthereum::mainnet();
+
+            let block = provider
+                .get_block(BlockId::number(block_number), BlockTransactionsKind::Full)
+                .await
+                .unwrap()
+                .unwrap();
+            let spec_id = chain.get_block_spec(&block.header).unwrap();
+            let rpc_storage =
+                pevm::RpcStorage::new(provider, spec_id, BlockId::number(block_number - 1));
+            common::test_execute_alloy(&chain, &rpc_storage, block, true);
+        }
+        #[cfg(feature = "optimism")]        
+        Chain::Op => {
+            use pevm::chain::PevmOptimism;
+            let provider = ProviderBuilder::new()
+                .network::<op_alloy_network::Optimism>()
+                .on_http(rpc_url.clone());
+            let chain = PevmOptimism::mainnet();
+
+            let block = provider
+                .get_block(BlockId::number(block_number), BlockTransactionsKind::Full)
+                .await
+                .unwrap()
+                .unwrap();
+            let spec_id = chain.get_block_spec(&block.header).unwrap();
+            let rpc_storage =
+                pevm::RpcStorage::new(provider, spec_id, BlockId::number(block_number - 1));
+            common::test_execute_alloy(&chain, &rpc_storage, block, true);
+        }
+    };
+}
+
+#[tokio::test(flavor = "multi_thread")]
+#[cfg(feature = "rpc-storage")]
+async fn mainnet_blocks_from_rpc() {
 
     let rpc_url = match std::env::var("ETHEREUM_RPC_URL") {
         // The empty check is for GitHub Actions where the variable is set with an empty string when unset!?
@@ -33,22 +77,16 @@ async fn mainnet_blocks_from_rpc() {
                // 17035010, // SHANGHAI
                // 19426587, // CANCUN
     ] {
-        let provider = ProviderBuilder::new().on_http(rpc_url.clone());
-        let block = provider
-            .get_block(BlockId::number(block_number), BlockTransactionsKind::Full)
-            .await
-            .unwrap()
-            .unwrap();
-        let chain = PevmEthereum::mainnet();
-        let spec_id = chain.get_block_spec(&block.header).unwrap();
-        let rpc_storage =
-            pevm::RpcStorage::new(provider, spec_id, BlockId::number(block_number - 1));
-        common::test_execute_alloy(&chain, &rpc_storage, block, true);
+      
+        blocks_from_rpc(&rpc_url,block_number,Chain::Op).await;
+
     }
 }
 
 #[test]
 fn mainnet_blocks_from_disk() {
+    use pevm::chain::PevmEthereum;
+
     common::for_each_block_from_disk(|block, storage| {
         // Run several times to try catching a race condition if there is any.
         // 1000~2000 is a better choice for local testing after major changes.
@@ -63,7 +101,6 @@ fn mainnet_blocks_from_disk() {
 async fn optimism_mainnet_blocks_from_rpc() {
     use alloy_provider::{Provider, ProviderBuilder};
     use alloy_rpc_types_eth::{BlockId, BlockTransactionsKind};
-    use pevm::chain::{PevmChain, PevmOptimism};
 
     let rpc_url = match std::env::var("OPTIMISM_RPC_URL") {
         Ok(value) if !value.is_empty() => value.parse().unwrap(),
@@ -78,20 +115,6 @@ async fn optimism_mainnet_blocks_from_rpc() {
                   // 122874325, // FJORD (https://specs.optimism.io/protocol/fjord/overview.html)
                   // 125874340, // GRANITE (https://specs.optimism.io/protocol/granite/overview.html)
     ] {
-        let provider = ProviderBuilder::new()
-            .network::<op_alloy_network::Optimism>()
-            .on_http(rpc_url.clone());
-        let block = provider
-            .get_block(BlockId::number(block_number), BlockTransactionsKind::Full)
-            .await
-            .unwrap()
-            .unwrap();
-
-        let chain = PevmOptimism::mainnet();
-        let spec_id = chain.get_block_spec(&block.header).unwrap();
-
-        let rpc_storage =
-            pevm::RpcStorage::new(provider, spec_id, BlockId::number(block_number - 1));
-        common::test_execute_alloy(&chain, &rpc_storage, block, true);
+        blocks_from_rpc(&rpc_url,block_number,Chain::Op).await;
     }
 }