diff --git a/crates/stages/src/stages/index_account_history.rs b/crates/stages/src/stages/index_account_history.rs index 2062b4eb02..fe0b6d3b40 100644 --- a/crates/stages/src/stages/index_account_history.rs +++ b/crates/stages/src/stages/index_account_history.rs @@ -75,17 +75,26 @@ mod tests { use std::collections::BTreeMap; use super::*; - use crate::test_utils::TestTransaction; + use crate::test_utils::{ + stage_test_suite_ext, ExecuteStageTestRunner, StageTestRunner, TestRunnerError, + TestTransaction, UnwindStageTestRunner, + }; + use itertools::Itertools; use reth_db::{ + cursor::DbCursorRO, models::{ - sharded_key::NUM_OF_INDICES_IN_SHARD, AccountBeforeTx, ShardedKey, + sharded_key, sharded_key::NUM_OF_INDICES_IN_SHARD, AccountBeforeTx, ShardedKey, StoredBlockBodyIndices, }, tables, - transaction::DbTxMut, + transaction::{DbTx, DbTxMut}, BlockNumberList, }; - use reth_primitives::{hex_literal::hex, H160, MAINNET}; + use reth_interfaces::test_utils::{ + generators, + generators::{random_block_range, random_contract_account_range, random_transition_range}, + }; + use reth_primitives::{hex_literal::hex, Address, BlockNumber, H160, H256, MAINNET}; const ADDRESS: H160 = H160(hex!("0000000000000000000000000000000000000001")); @@ -357,4 +366,137 @@ mod tests { ]) ); } + + stage_test_suite_ext!(IndexAccountHistoryTestRunner, index_account_history); + + struct IndexAccountHistoryTestRunner { + pub(crate) tx: TestTransaction, + commit_threshold: u64, + } + + impl Default for IndexAccountHistoryTestRunner { + fn default() -> Self { + Self { tx: TestTransaction::default(), commit_threshold: 1000 } + } + } + + impl StageTestRunner for IndexAccountHistoryTestRunner { + type S = IndexAccountHistoryStage; + + fn tx(&self) -> &TestTransaction { + &self.tx + } + + fn stage(&self) -> Self::S { + Self::S { commit_threshold: self.commit_threshold } + } + } + + impl ExecuteStageTestRunner for IndexAccountHistoryTestRunner { + type Seed = (); + + fn seed_execution(&mut self, input: ExecInput) -> Result { + let stage_process = input.checkpoint().block_number; + let start = stage_process + 1; + let end = input.target(); + let mut rng = generators::rng(); + + let num_of_accounts = 31; + let accounts = random_contract_account_range(&mut rng, &mut (0..num_of_accounts)) + .into_iter() + .collect::>(); + + let blocks = random_block_range(&mut rng, start..=end, H256::zero(), 0..3); + + let (transitions, _) = random_transition_range( + &mut rng, + blocks.iter(), + accounts.into_iter().map(|(addr, acc)| (addr, (acc, Vec::new()))), + 0..3, + 0..256, + ); + + // add block changeset from block 1. + self.tx.insert_transitions(transitions, Some(start))?; + + Ok(()) + } + + fn validate_execution( + &self, + input: ExecInput, + output: Option, + ) -> Result<(), TestRunnerError> { + if let Some(output) = output { + let start_block = input.next_block(); + let end_block = output.checkpoint.block_number; + if start_block > end_block { + return Ok(()) + } + + assert_eq!( + output, + ExecOutput { checkpoint: StageCheckpoint::new(input.target()), done: true } + ); + + let provider = self.tx.inner(); + let mut changeset_cursor = + provider.tx_ref().cursor_read::()?; + + let account_transitions = + changeset_cursor.walk_range(start_block..=end_block)?.try_fold( + BTreeMap::new(), + |mut accounts: BTreeMap>, + entry| + -> Result<_, TestRunnerError> { + let (index, account) = entry?; + accounts.entry(account.address).or_default().push(index); + Ok(accounts) + }, + )?; + + let mut result = BTreeMap::new(); + for (address, indices) in account_transitions { + // chunk indices and insert them in shards of N size. + let mut chunks = indices + .iter() + .chunks(sharded_key::NUM_OF_INDICES_IN_SHARD) + .into_iter() + .map(|chunks| chunks.map(|i| *i as usize).collect::>()) + .collect::>(); + let last_chunk = chunks.pop(); + + chunks.into_iter().for_each(|list| { + result.insert( + ShardedKey::new( + address, + *list.last().expect("Chuck does not return empty list") + as BlockNumber, + ) as ShardedKey, + list, + ); + }); + + if let Some(last_list) = last_chunk { + result.insert( + ShardedKey::new(address, u64::MAX) as ShardedKey, + last_list, + ); + }; + } + + let table = cast(self.tx.table::().unwrap()); + assert_eq!(table, result); + } + Ok(()) + } + } + + impl UnwindStageTestRunner for IndexAccountHistoryTestRunner { + fn validate_unwind(&self, _input: UnwindInput) -> Result<(), TestRunnerError> { + let table = self.tx.table::().unwrap(); + assert!(table.is_empty()); + Ok(()) + } + } }