Modularize mock crate for easier and more customizable testing setup generation

bus-mapping/src/evm/opcodes/calldatacopy.rs

@@ -157,11 +157,19 @@ fn gen_memory_copy_steps(
 
 #[cfg(test)]
 mod calldatacopy_tests {
-    use super::*;
-    use crate::circuit_input_builder::ExecState;
-    use crate::operation::StackOp;
-    use eth_types::evm_types::{OpcodeId, StackAddress};
-    use eth_types::{bytecode, Word};
+    use crate::{
+        circuit_input_builder::ExecState,
+        mock::BlockData,
+        operation::{CallContextField, CallContextOp, StackOp, RW},
+    };
+    use eth_types::{
+        bytecode,
+        evm_types::{OpcodeId, StackAddress},
+        geth_types::GethData,
+        Word,
+    };
+
+    use mock::test_ctx::{helpers::*, TestContext};
     use pretty_assertions::assert_eq;
 
     #[test]
@@ -175,11 +183,16 @@ mod calldatacopy_tests {
         };
 
         // Get the execution steps from the external tracer
-        let block = crate::mock::BlockData::new_from_geth_data(
-            mock::new_single_tx_trace_code(&code).unwrap(),
-        );
-
-        let mut builder = block.new_circuit_input_builder();
+        let block: GethData = TestContext::<2, 1>::new(
+            None,
+            account_0_code_account_1_no_code(code),
+            tx_from_1_to_0,
+            |block, _tx| block.number(0xcafeu64),
+        )
+        .unwrap()
+        .into();
+
+        let mut builder = BlockData::new_from_geth_data(block.clone()).new_circuit_input_builder();
         builder
             .handle_block(&block.eth_block, &block.geth_traces)
             .unwrap();

bus-mapping/src/evm/opcodes/calldatasize.rs

@@ -40,9 +40,18 @@ impl Opcode for Calldatasize {
 
 #[cfg(test)]
 mod calldatasize_tests {
-    use crate::circuit_input_builder::ExecState;
-    use crate::operation::{CallContextField, CallContextOp, StackOp, RW};
-    use eth_types::{bytecode, evm_types::OpcodeId, evm_types::StackAddress};
+    use crate::{
+        circuit_input_builder::ExecState,
+        mock::BlockData,
+        operation::{CallContextField, CallContextOp, StackOp, RW},
+    };
+    use eth_types::{
+        bytecode,
+        evm_types::{OpcodeId, StackAddress},
+        geth_types::GethData,
+    };
+
+    use mock::test_ctx::{helpers::*, TestContext};
     use pretty_assertions::assert_eq;
 
     #[test]
@@ -53,11 +62,16 @@ mod calldatasize_tests {
         };
 
         // Get the execution steps from the external tracer
-        let block = crate::mock::BlockData::new_from_geth_data(
-            mock::new_single_tx_trace_code(&code).unwrap(),
-        );
+        let block: GethData = TestContext::<2, 1>::new(
+            None,
+            account_0_code_account_1_no_code(code),
+            tx_from_1_to_0,
+            |block, _tx| block.number(0xcafeu64),
+        )
+        .unwrap()
+        .into();
 
-        let mut builder = block.new_circuit_input_builder();
+        let mut builder = BlockData::new_from_geth_data(block.clone()).new_circuit_input_builder();
         builder
             .handle_block(&block.eth_block, &block.geth_traces)
             .unwrap();

bus-mapping/src/evm/opcodes/caller.rs

@@ -42,9 +42,16 @@ impl Opcode for Caller {
 
 #[cfg(test)]
 mod caller_tests {
-    use crate::circuit_input_builder::ExecState;
-    use crate::operation::{CallContextField, CallContextOp, StackOp, RW};
-    use eth_types::{bytecode, evm_types::OpcodeId, evm_types::StackAddress, ToWord};
+    use super::*;
+    use crate::{circuit_input_builder::ExecState, mock::BlockData, operation::StackOp};
+    use eth_types::{
+        bytecode,
+        evm_types::{OpcodeId, StackAddress},
+        geth_types::GethData,
+        ToWord,
+    };
+
+    use mock::test_ctx::{helpers::*, TestContext};
     use pretty_assertions::assert_eq;
 
     #[test]
@@ -55,11 +62,16 @@ mod caller_tests {
         };
 
         // Get the execution steps from the external tracer
-        let block = crate::mock::BlockData::new_from_geth_data(
-            mock::new_single_tx_trace_code(&code).unwrap(),
-        );
+        let block: GethData = TestContext::<2, 1>::new(
+            None,
+            account_0_code_account_1_no_code(code),
+            tx_from_1_to_0,
+            |block, _tx| block.number(0xcafeu64),
+        )
+        .unwrap()
+        .into();
 
-        let mut builder = block.new_circuit_input_builder();
+        let mut builder = BlockData::new_from_geth_data(block.clone()).new_circuit_input_builder();
         builder
             .handle_block(&block.eth_block, &block.geth_traces)
             .unwrap();

bus-mapping/src/evm/opcodes/callvalue.rs

@@ -42,9 +42,17 @@ impl Opcode for Callvalue {
 
 #[cfg(test)]
 mod callvalue_tests {
-    use crate::circuit_input_builder::ExecState;
-    use crate::operation::{CallContextField, CallContextOp, StackOp, RW};
-    use eth_types::{bytecode, evm_types::OpcodeId, evm_types::StackAddress};
+    use crate::{
+        circuit_input_builder::ExecState,
+        mock::BlockData,
+        operation::{CallContextField, CallContextOp, StackOp, RW},
+    };
+    use eth_types::{
+        bytecode,
+        evm_types::{OpcodeId, StackAddress},
+        geth_types::GethData,
+    };
+    use mock::test_ctx::{helpers::*, TestContext};
     use pretty_assertions::assert_eq;
 
     #[test]
@@ -55,11 +63,16 @@ mod callvalue_tests {
         };
 
         // Get the execution steps from the external tracer
-        let block = crate::mock::BlockData::new_from_geth_data(
-            mock::new_single_tx_trace_code(&code).unwrap(),
-        );
+        let block: GethData = TestContext::<2, 1>::new(
+            None,
+            account_0_code_account_1_no_code(code),
+            tx_from_1_to_0,
+            |block, _tx| block.number(0xcafeu64),
+        )
+        .unwrap()
+        .into();
 
-        let mut builder = block.new_circuit_input_builder();
+        let mut builder = BlockData::new_from_geth_data(block.clone()).new_circuit_input_builder();
         builder
             .handle_block(&block.eth_block, &block.geth_traces)
             .unwrap();

bus-mapping/src/evm/opcodes/dup.rs

@@ -34,11 +34,10 @@ impl<const N: usize> Opcode for Dup<N> {
 #[cfg(test)]
 mod dup_tests {
     use super::*;
-    use crate::operation::StackOp;
-    use eth_types::bytecode;
-    use eth_types::evm_types::StackAddress;
-    use eth_types::word;
+    use crate::{mock::BlockData, operation::StackOp};
+    use eth_types::{bytecode, evm_types::StackAddress, geth_types::GethData, word};
     use itertools::Itertools;
+    use mock::test_ctx::{helpers::*, TestContext};
     use pretty_assertions::assert_eq;
 
     #[test]
@@ -54,11 +53,16 @@ mod dup_tests {
         };
 
         // Get the execution steps from the external tracer
-        let block = crate::mock::BlockData::new_from_geth_data(
-            mock::new_single_tx_trace_code(&code).unwrap(),
-        );
+        let block: GethData = TestContext::<2, 1>::new(
+            None,
+            account_0_code_account_1_no_code(code),
+            tx_from_1_to_0,
+            |block, _tx| block.number(0xcafeu64),
+        )
+        .unwrap()
+        .into();
 
-        let mut builder = block.new_circuit_input_builder();
+        let mut builder = BlockData::new_from_geth_data(block.clone()).new_circuit_input_builder();
         builder
             .handle_block(&block.eth_block, &block.geth_traces)
             .unwrap();
@@ -73,7 +77,6 @@ mod dup_tests {
                 .iter()
                 .filter(|step| step.exec_state.is_dup())
                 .collect_vec()[i];
-
             assert_eq!(
                 [0, 1]
                     .map(|idx| &builder.block.container.stack

bus-mapping/src/evm/opcodes/extcodehash.rs

@@ -133,11 +133,13 @@ mod extcodehash_tests {
     use crate::mock::BlockData;
     use crate::operation::StackOp;
     use eth_types::{
-        address, bytecode, evm_types::OpcodeId, geth_types::Account as GethAccount, Address,
-        Bytecode, Bytes, U256,
+        address, bytecode,
+        evm_types::{OpcodeId, StackAddress},
+        geth_types::GethData,
+        Bytecode, Bytes, Word, U256,
     };
     use ethers_core::utils::keccak256;
-    use mock::new_single_tx_trace_accounts;
+    use mock::TestContext;
     use pretty_assertions::assert_eq;
 
     #[test]
@@ -178,36 +180,49 @@ mod extcodehash_tests {
             EXTCODEHASH
             STOP
         });
+        let mut nonce = Word::from(300u64);
+        let mut balance = Word::from(800u64);
+        let mut code_ext = Bytes::from([34, 54, 56]);
 
-        let mut accounts = vec![GethAccount {
-            address: Address::default(), // This is the address of the executing account
-            code: Bytes::from(code.to_vec()),
-            ..Default::default()
-        }];
-
-        // Let the external account exist, if needed, by making its code non-empty
-        let (mut nonce, mut balance, mut code) = (U256::zero(), U256::zero(), Bytes::default());
-        if exists {
-            nonce = U256::from(300u64);
-            balance = U256::from(800u64);
-            code = Bytes::from([34, 54, 56]);
-
-            accounts.push(GethAccount {
-                address: external_address,
-                nonce,
-                balance,
-                code: code.clone(),
-                ..Default::default()
-            })
+        if !exists {
+            nonce = Word::zero();
+            balance = Word::zero();
+            code_ext = Bytes::default();
         }
-        let code_hash = keccak256(&code).into();
 
         // Get the execution steps from the external tracer
-        let geth_data = new_single_tx_trace_accounts(accounts)?;
-        let block = BlockData::new_from_geth_data(geth_data);
+        let block: GethData = TestContext::<3, 1>::new(
+            None,
+            |accs| {
+                accs[0]
+                    .address(address!("0x0000000000000000000000000000000000000010"))
+                    .balance(Word::from(1u64 << 20))
+                    .code(code.clone());
+
+                accs[1]
+                    .address(external_address)
+                    .balance(balance)
+                    .nonce(nonce)
+                    .code(code_ext.clone());
+
+                accs[2]
+                    .address(address!("0x0000000000000000000000000000000000cafe01"))
+                    .balance(Word::from(1u64 << 20));
+            },
+            |mut txs, accs| {
+                txs[0].to(accs[0].address).from(accs[2].address);
+            },
+            |block, _tx| block.number(0xcafeu64),
+        )
+        .unwrap()
+        .into();
+
+        let code_hash = Word::from(keccak256(code_ext));
 
-        let mut builder = block.new_circuit_input_builder();
-        builder.handle_block(&block.eth_block, &block.geth_traces)?;
+        let mut builder = BlockData::new_from_geth_data(block.clone()).new_circuit_input_builder();
+        builder
+            .handle_block(&block.eth_block, &block.geth_traces)
+            .unwrap();
 
         // Check that `external_address` is in access list as a result of bus mapping.
         assert!(builder.sdb.add_account_to_access_list(external_address));
@@ -233,7 +248,7 @@ mod extcodehash_tests {
                 RW::READ,
                 &StackOp {
                     call_id,
-                    address: 1023.into(),
+                    address: StackAddress::from(1023u32),
                     value: external_address.to_word()
                 }
             )
@@ -349,7 +364,7 @@ mod extcodehash_tests {
                 RW::WRITE,
                 &StackOp {
                     call_id,
-                    address: 1023.into(),
+                    address: 1023u32.into(),
                     value: if exists { code_hash } else { U256::zero() }
                 }
             )

bus-mapping/src/evm/opcodes/mload.rs

@@ -60,12 +60,18 @@ impl Opcode for Mload {
 #[cfg(test)]
 mod mload_tests {
     use super::*;
-    use crate::circuit_input_builder::ExecState;
-    use crate::operation::{MemoryOp, StackOp};
-    use eth_types::bytecode;
-    use eth_types::evm_types::{OpcodeId, StackAddress};
-    use eth_types::Word;
+    use crate::{
+        circuit_input_builder::ExecState,
+        mock::BlockData,
+        operation::{MemoryOp, StackOp},
+    };
+    use eth_types::{
+        bytecode,
+        evm_types::{OpcodeId, StackAddress},
+        geth_types::GethData,
+    };
     use itertools::Itertools;
+    use mock::test_ctx::{helpers::*, TestContext};
     use pretty_assertions::assert_eq;
 
     #[test]
@@ -79,11 +85,16 @@ mod mload_tests {
         };
 
         // Get the execution steps from the external tracer
-        let block = crate::mock::BlockData::new_from_geth_data(
-            mock::new_single_tx_trace_code(&code).unwrap(),
-        );
+        let block: GethData = TestContext::<2, 1>::new(
+            None,
+            account_0_code_account_1_no_code(code),
+            tx_from_1_to_0,
+            |block, _tx| block.number(0xcafeu64),
+        )
+        .unwrap()
+        .into();
 
-        let mut builder = block.new_circuit_input_builder();
+        let mut builder = BlockData::new_from_geth_data(block.clone()).new_circuit_input_builder();
         builder
             .handle_block(&block.eth_block, &block.geth_traces)
             .unwrap();