test: add more tests to the DB code we wrote (#144)

* test: add more tests to the DB code we wrote

* chore: hide

Author: prestwich

src/db/cow/mod.rs

@@ -258,7 +258,9 @@ impl<Db> DatabaseCommit for CacheOnWrite<Db> {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::db::test_utils::DbTestExt;
     use revm::{
+        bytecode::Bytecode,
         database::InMemoryDB,
         state::{AccountStatus, EvmStorageSlot},
     };
@@ -320,6 +322,100 @@ mod tests {
         assert_eq!(cow_db.storage(address, U256::from(0)).unwrap(), U256::from(42));
         assert_eq!(cow_db.storage(address, U256::from(1)).unwrap(), U256::from(42));
     }
+
+    #[test]
+    fn test_set_balance() {
+        let addr = Address::repeat_byte(1);
+        let mut cow = CacheOnWrite::new(InMemoryDB::default());
+
+        cow.test_set_balance(addr, U256::from(1000));
+        assert_eq!(cow.basic(addr).unwrap().unwrap().balance, U256::from(1000));
+
+        cow.test_increase_balance(addr, U256::from(500));
+        assert_eq!(cow.basic(addr).unwrap().unwrap().balance, U256::from(1500));
+
+        cow.test_decrease_balance(addr, U256::from(200));
+        assert_eq!(cow.basic(addr).unwrap().unwrap().balance, U256::from(1300));
+    }
+
+    #[test]
+    fn test_set_nonce() {
+        let addr = Address::repeat_byte(2);
+        let mut cow = CacheOnWrite::new(InMemoryDB::default());
+
+        cow.test_set_nonce(addr, 42);
+        assert_eq!(cow.basic(addr).unwrap().unwrap().nonce, 42);
+    }
+
+    #[test]
+    fn test_set_storage() {
+        let addr = Address::repeat_byte(3);
+        let mut cow = CacheOnWrite::new(InMemoryDB::default());
+
+        cow.test_set_storage(addr, U256::from(10), U256::from(999));
+        assert_eq!(cow.storage(addr, U256::from(10)).unwrap(), U256::from(999));
+
+        // Set another slot, verify first slot still readable
+        cow.test_set_storage(addr, U256::from(20), U256::from(888));
+        assert_eq!(cow.storage(addr, U256::from(10)).unwrap(), U256::from(999));
+        assert_eq!(cow.storage(addr, U256::from(20)).unwrap(), U256::from(888));
+    }
+
+    #[test]
+    fn test_set_bytecode() {
+        let addr = Address::repeat_byte(4);
+        let mut cow = CacheOnWrite::new(InMemoryDB::default());
+        let code =
+            Bytecode::new_raw(alloy::primitives::Bytes::from_static(&[0x60, 0x00, 0x60, 0x00]));
+        let code_hash = code.hash_slow();
+
+        cow.test_set_bytecode(addr, code.clone());
+
+        let info = cow.basic(addr).unwrap().unwrap();
+        assert_eq!(info.code_hash, code_hash);
+        assert_eq!(cow.code_by_hash(code_hash).unwrap(), code);
+    }
+
+    #[test]
+    fn test_cow_isolates_writes_from_inner() {
+        let addr = Address::repeat_byte(5);
+        let mut inner = InMemoryDB::default();
+        inner.test_set_balance(addr, U256::from(100));
+
+        let mut cow = CacheOnWrite::new(inner);
+        cow.test_set_balance(addr, U256::from(500));
+
+        // COW sees updated balance
+        assert_eq!(cow.basic(addr).unwrap().unwrap().balance, U256::from(500));
+        // Inner still has original via ref
+        assert_eq!(cow.inner().basic_ref(addr).unwrap().unwrap().balance, U256::from(100));
+    }
+
+    #[test]
+    fn test_nested_cow_flatten() {
+        let addr1 = Address::repeat_byte(6);
+        let addr2 = Address::repeat_byte(7);
+        let mut inner = InMemoryDB::default();
+        inner.test_set_balance(addr1, U256::from(100));
+
+        let mut cow = CacheOnWrite::new(inner);
+        cow.test_set_balance(addr2, U256::from(50));
+
+        let mut nested = cow.nest();
+        nested.test_set_balance(addr1, U256::from(200));
+
+        // Flatten nested -> cow
+        let mut cow = nested.flatten();
+        // addr1 was overwritten in nested layer
+        assert_eq!(cow.basic(addr1).unwrap().unwrap().balance, U256::from(200));
+        // addr2 untouched in nested, visible from cow layer
+        assert_eq!(cow.basic(addr2).unwrap().unwrap().balance, U256::from(50));
+
+        // Flatten cow -> inner
+        let inner = cow.flatten();
+        assert_eq!(inner.basic_ref(addr1).unwrap().unwrap().balance, U256::from(200));
+        assert_eq!(inner.basic_ref(addr2).unwrap().unwrap().balance, U256::from(50));
+    }
 }
 
 // Some code above and documentation is adapted from the revm crate, and is

src/db/mod.rs

@@ -13,3 +13,7 @@ pub mod cow;
 #[cfg(feature = "alloy-db")]
 /// Alloy-powered revm Database implementation that fetches data over the network.
 pub mod alloy;
+
+/// Test utilities for databases.
+#[doc(hidden)]
+pub mod test_utils;

src/db/sync/state.rs

@@ -477,6 +477,89 @@ mod test {
             U256::from(100)
         );
     }
+
+    #[test]
+    fn test_increment_balances() {
+        let addr = Address::repeat_byte(10);
+        let mut state = ConcurrentState::new(EmptyDB::new(), Default::default());
+
+        state.increment_balances([(addr, 1000)]).unwrap();
+        assert_eq!(state.basic(addr).unwrap().unwrap().balance, U256::from(1000));
+
+        state.increment_balances([(addr, 500)]).unwrap();
+        assert_eq!(state.basic(addr).unwrap().unwrap().balance, U256::from(1500));
+    }
+
+    #[test]
+    fn test_drain_balances() {
+        let addr = Address::repeat_byte(11);
+        let mut state = ConcurrentState::new(EmptyDB::new(), Default::default());
+
+        state.increment_balances([(addr, 1000)]).unwrap();
+        let drained = state.drain_balances([addr]).unwrap();
+        assert_eq!(drained, vec![1000]);
+        assert_eq!(state.basic(addr).unwrap().unwrap().balance, U256::ZERO);
+    }
+
+    #[test]
+    fn test_insert_account() {
+        let addr = Address::repeat_byte(12);
+        let mut state = ConcurrentState::new(EmptyDB::new(), Default::default());
+
+        let info = AccountInfo { balance: U256::from(999), nonce: 42, ..Default::default() };
+        state.insert_account(addr, info);
+
+        let loaded = state.basic(addr).unwrap().unwrap();
+        assert_eq!(loaded.balance, U256::from(999));
+        assert_eq!(loaded.nonce, 42);
+    }
+
+    #[test]
+    fn test_child_isolation() {
+        let addr = Address::repeat_byte(14);
+        let parent = Arc::new(ConcurrentState::new(EmptyDB::new(), Default::default()));
+
+        // Use increment_balances which properly handles account loading
+        let mut child = parent.child();
+        child.increment_balances([(addr, 500)]).unwrap();
+
+        // Child sees balance
+        assert_eq!(
+            child.info.cache.accounts.get(&addr).unwrap().account_info().unwrap().balance,
+            U256::from(500)
+        );
+        // Parent sees account loaded but with zero balance (from EmptyDB lookup)
+        assert_eq!(
+            parent.info.cache.accounts.get(&addr).unwrap().account_info(),
+            None // Not existing account loaded
+        );
+
+        // After merge, parent sees child's value
+        let mut parent = parent;
+        parent.merge_child(child).unwrap();
+        assert_eq!(parent.basic_ref(addr).unwrap().unwrap().balance, U256::from(500));
+    }
+
+    #[test]
+    fn test_child_merge_accumulates() {
+        let addr = Address::repeat_byte(15);
+        let mut parent = Arc::new(ConcurrentState::new(EmptyDB::new(), Default::default()));
+
+        // Set initial balance on parent
+        Arc::get_mut(&mut parent).unwrap().increment_balances([(addr, 100)]).unwrap();
+
+        let mut child = parent.child();
+        child.increment_balances([(addr, 50)]).unwrap();
+
+        // Child sees accumulated balance
+        assert_eq!(child.basic(addr).unwrap().unwrap().balance, U256::from(150));
+        // Parent sees original
+        assert_eq!(parent.basic_ref(addr).unwrap().unwrap().balance, U256::from(100));
+
+        // After merge
+        parent.merge_child(child).unwrap();
+        assert_eq!(parent.basic_ref(addr).unwrap().unwrap().balance, U256::from(150));
+    }
 }
 
 // Some code above and documentation is adapted from the revm crate, and is

src/db/test_utils.rs

@@ -0,0 +1,162 @@
+use alloy::primitives::{Address, U256};
+use revm::{
+    bytecode::Bytecode,
+    database::{Database, DatabaseCommit, TryDatabaseCommit},
+    primitives::HashMap,
+    state::{Account, AccountInfo, AccountStatus, EvmStorageSlot},
+};
+
+/// Test utilities for databases.
+#[doc(hidden)]
+pub trait DbTestExt: Database + DatabaseCommit {
+    /// Modify an account via closure, creating it if it doesn't exist.
+    fn test_modify_account<F>(&mut self, address: Address, f: F)
+    where
+        F: FnOnce(&mut AccountInfo),
+    {
+        let mut info = self.basic(address).ok().flatten().unwrap_or_default();
+        f(&mut info);
+        let account = Account {
+            info,
+            storage: Default::default(),
+            status: AccountStatus::Touched,
+            transaction_id: 0,
+        };
+        self.commit(HashMap::from_iter([(address, account)]));
+    }
+
+    /// Set account balance.
+    fn test_set_balance(&mut self, address: Address, balance: U256) {
+        self.test_modify_account(address, |info| info.balance = balance);
+    }
+
+    /// Increase account balance (saturating).
+    fn test_increase_balance(&mut self, address: Address, amount: U256) {
+        self.test_modify_account(address, |info| {
+            info.balance = info.balance.saturating_add(amount)
+        });
+    }
+
+    /// Decrease account balance (saturating).
+    fn test_decrease_balance(&mut self, address: Address, amount: U256) {
+        self.test_modify_account(address, |info| {
+            info.balance = info.balance.saturating_sub(amount)
+        });
+    }
+
+    /// Set account nonce.
+    fn test_set_nonce(&mut self, address: Address, nonce: u64) {
+        self.test_modify_account(address, |info| info.nonce = nonce);
+    }
+
+    /// Set a storage slot.
+    fn test_set_storage(&mut self, address: Address, slot: U256, value: U256) {
+        let info = self.basic(address).ok().flatten().unwrap_or_default();
+        let storage =
+            HashMap::from_iter([(slot, EvmStorageSlot::new_changed(U256::ZERO, value, 0))]);
+        let account = Account { info, storage, status: AccountStatus::Touched, transaction_id: 0 };
+        self.commit(HashMap::from_iter([(address, account)]));
+    }
+
+    /// Set account bytecode.
+    fn test_set_bytecode(&mut self, address: Address, bytecode: Bytecode) {
+        self.test_modify_account(address, |info| {
+            info.code_hash = bytecode.hash_slow();
+            info.code = Some(bytecode);
+        });
+    }
+}
+
+impl<Db: Database + DatabaseCommit> DbTestExt for Db {}
+
+/// Test utilities for databases that support fallible commits.
+#[doc(hidden)]
+pub trait TryDbTestExt: Database + TryDatabaseCommit {
+    /// Modify an account via closure, creating it if it doesn't exist.
+    fn test_try_modify_account<F>(
+        &mut self,
+        address: Address,
+        f: F,
+    ) -> Result<(), <Self as TryDatabaseCommit>::Error>
+    where
+        F: FnOnce(&mut AccountInfo),
+    {
+        let mut info = self.basic(address).ok().flatten().unwrap_or_default();
+        f(&mut info);
+        let account = Account {
+            info,
+            storage: Default::default(),
+            status: AccountStatus::Touched,
+            transaction_id: 0,
+        };
+        self.try_commit(HashMap::from_iter([(address, account)]))
+    }
+
+    /// Set account balance.
+    fn test_try_set_balance(
+        &mut self,
+        address: Address,
+        balance: U256,
+    ) -> Result<(), <Self as TryDatabaseCommit>::Error> {
+        self.test_try_modify_account(address, |info| info.balance = balance)
+    }
+
+    /// Increase account balance (saturating).
+    fn test_try_increase_balance(
+        &mut self,
+        address: Address,
+        amount: U256,
+    ) -> Result<(), <Self as TryDatabaseCommit>::Error> {
+        self.test_try_modify_account(address, |info| {
+            info.balance = info.balance.saturating_add(amount)
+        })
+    }
+
+    /// Decrease account balance (saturating).
+    fn test_try_decrease_balance(
+        &mut self,
+        address: Address,
+        amount: U256,
+    ) -> Result<(), <Self as TryDatabaseCommit>::Error> {
+        self.test_try_modify_account(address, |info| {
+            info.balance = info.balance.saturating_sub(amount)
+        })
+    }
+
+    /// Set account nonce.
+    fn test_try_set_nonce(
+        &mut self,
+        address: Address,
+        nonce: u64,
+    ) -> Result<(), <Self as TryDatabaseCommit>::Error> {
+        self.test_try_modify_account(address, |info| info.nonce = nonce)
+    }
+
+    /// Set a storage slot.
+    fn test_try_set_storage(
+        &mut self,
+        address: Address,
+        slot: U256,
+        value: U256,
+    ) -> Result<(), <Self as TryDatabaseCommit>::Error> {
+        let info = self.basic(address).ok().flatten().unwrap_or_default();
+        let storage =
+            HashMap::from_iter([(slot, EvmStorageSlot::new_changed(U256::ZERO, value, 0))]);
+        let account = Account { info, storage, status: AccountStatus::Touched, transaction_id: 0 };
+        self.try_commit(HashMap::from_iter([(address, account)]))
+    }
+
+    /// Set account bytecode.
+    fn test_try_set_bytecode(
+        &mut self,
+        address: Address,
+        bytecode: Bytecode,
+    ) -> Result<(), <Self as TryDatabaseCommit>::Error> {
+        self.test_try_modify_account(address, |info| {
+            info.code_hash = bytecode.hash_slow();
+            info.code = Some(bytecode);
+        })
+    }
+}
+
+impl<Db: Database + TryDatabaseCommit> TryDbTestExt for Db {}