From eea3d616f75f995f4f4d17793b39fa02ed7e3446 Mon Sep 17 00:00:00 2001 From: Tayfun Elmas Date: Tue, 5 Nov 2024 17:43:31 +0300 Subject: [PATCH] feat(contract-distribution): Add few unittests for ContractStorage (#12380) --- core/store/src/contract.rs | 211 +++++++++++++++++++++++++++++++++++++ 1 file changed, 211 insertions(+) diff --git a/core/store/src/contract.rs b/core/store/src/contract.rs index 60393990b82..d79ab1c5c44 100644 --- a/core/store/src/contract.rs +++ b/core/store/src/contract.rs @@ -180,3 +180,214 @@ impl ContractStorage { tracker.finalize() } } + +#[cfg(test)] +mod tests { + use std::{ + collections::{HashMap, HashSet}, + sync::Arc, + }; + + use itertools::Itertools; + use near_primitives::{ + errors::{MissingTrieValueContext, StorageError}, + hash::CryptoHash, + stateless_validation::contract_distribution::CodeHash, + }; + use near_vm_runner::ContractCode; + + use crate::{contract::ContractStorage, TrieStorage}; + + struct MockTrieStorage { + store: HashMap>, + } + + impl MockTrieStorage { + fn new() -> Self { + Self { store: HashMap::new() } + } + + fn insert(&mut self, hash: CryptoHash, data: Arc<[u8]>) { + self.store.insert(hash, data); + } + } + + impl TrieStorage for MockTrieStorage { + fn retrieve_raw_bytes(&self, hash: &CryptoHash) -> Result, StorageError> { + match self.store.get(hash) { + Some(data) => Ok(data.clone()), + None => { + Err(StorageError::MissingTrieValue(MissingTrieValueContext::TrieStorage, *hash)) + } + } + } + } + + /// Tests a scenario with old (already existing in the storage) and new contracts and finalizing after rolling back some deploys and committing others. + #[test] + fn test_contract_storage_finalize_after_rollback_and_commit() { + let contracts = (0..4).map(|i| ContractCode::new(vec![i], None)).collect_vec(); + let (old_contracts, new_contracts) = contracts.split_at(2); + + // Insert old contracts (already deployed contracts) into the storage. + let mut mock_storage = MockTrieStorage::new(); + for contract in old_contracts.iter() { + mock_storage.insert(*contract.hash(), contract.code().to_vec().into()); + } + + let mut contract_storage = ContractStorage::new(Arc::new(mock_storage)); + + contract_storage.record_deploy(old_contracts[0].clone_for_tests()); + contract_storage.record_deploy(new_contracts[0].clone_for_tests()); + + contract_storage.record_call(*old_contracts[0].hash()); + contract_storage.record_call(*new_contracts[0].hash()); + + contract_storage.rollback_deploys(); + + contract_storage.record_deploy(old_contracts[1].clone_for_tests()); + contract_storage.record_deploy(new_contracts[1].clone_for_tests()); + + contract_storage.record_call(*old_contracts[1].hash()); + contract_storage.record_call(*new_contracts[1].hash()); + + contract_storage.commit_deploys(); + + let updates = contract_storage.finalize(); + assert_eq!( + updates.contract_accesses, + HashSet::from_iter(vec![ + CodeHash(*old_contracts[0].hash()), + CodeHash(*old_contracts[1].hash()), + CodeHash(*new_contracts[0].hash()), + CodeHash(*new_contracts[1].hash()) + ]) + ); + assert_eq!( + updates.contract_deploy_hashes(), + HashSet::from_iter(vec![ + CodeHash(*old_contracts[1].hash()), + CodeHash(*new_contracts[1].hash()) + ]) + ); + } + + /// Tests a scenario with old (already existing in the storage) and new contracts and calling `get` after committing some deploys. + #[test] + fn test_contract_storage_get_after_new_deploys_and_commit() { + let contracts = (0..4).map(|i| ContractCode::new(vec![i], None)).collect_vec(); + let (old_contracts, new_contracts) = contracts.split_at(2); + + // Insert old contracts (already deployed contracts) into the storage. + let mut mock_storage = MockTrieStorage::new(); + for contract in old_contracts.iter() { + mock_storage.insert(*contract.hash(), contract.code().to_vec().into()); + } + + let mut contract_storage = ContractStorage::new(Arc::new(mock_storage)); + + // Only existing contracts should be returned by `get` before deploying the new ones. + for contract in old_contracts.iter() { + assert_eq!(contract_storage.get(*contract.hash()).unwrap().hash(), contract.hash()); + } + for contract in new_contracts.iter() { + assert!(contract_storage.get(*contract.hash()).is_none()); + } + + contract_storage.record_deploy(new_contracts[0].clone_for_tests()); + contract_storage.record_deploy(new_contracts[1].clone_for_tests()); + + // Make the same `get` calls before and after commit. Both should return the same results. + + for contract in old_contracts.iter() { + assert_eq!(contract_storage.get(*contract.hash()).unwrap().hash(), contract.hash()); + } + for contract in new_contracts.iter() { + assert_eq!(contract_storage.get(*contract.hash()).unwrap().hash(), contract.hash()); + } + + contract_storage.commit_deploys(); + + for contract in old_contracts.iter() { + assert_eq!(contract_storage.get(*contract.hash()).unwrap().hash(), contract.hash()); + } + for contract in new_contracts.iter() { + assert_eq!(contract_storage.get(*contract.hash()).unwrap().hash(), contract.hash()); + } + } + + /// Tests a scenario with old (already existing in the storage) and new contracts and calling `get` after rolling back some deploys. + #[test] + fn test_contract_storage_get_after_new_deploys_and_rollback() { + let contracts = (0..4).map(|i| ContractCode::new(vec![i], None)).collect_vec(); + let (old_contracts, new_contracts) = contracts.split_at(2); + + // Insert old contracts (already deployed contracts) into the storage. + let mut mock_storage = MockTrieStorage::new(); + for contract in old_contracts.iter() { + mock_storage.insert(*contract.hash(), contract.code().to_vec().into()); + } + + let mut contract_storage = ContractStorage::new(Arc::new(mock_storage)); + + contract_storage.record_deploy(new_contracts[0].clone_for_tests()); + contract_storage.record_deploy(new_contracts[1].clone_for_tests()); + + // Make the same `get` calls before and after rollback. Calls before rollback return the new contracts + // but those after rollback should return only the old contracts. + + for contract in old_contracts.iter() { + assert_eq!(contract_storage.get(*contract.hash()).unwrap().hash(), contract.hash()); + } + for contract in new_contracts.iter() { + assert_eq!(contract_storage.get(*contract.hash()).unwrap().hash(), contract.hash()); + } + + contract_storage.rollback_deploys(); + + for contract in old_contracts.iter() { + assert_eq!(contract_storage.get(*contract.hash()).unwrap().hash(), contract.hash()); + } + for contract in new_contracts.iter() { + assert!(contract_storage.get(*contract.hash()).is_none()); + } + } + + /// Tests a scenario with existing and missing contracts, and calling `get` after finalizing the storage. + #[test] + fn test_contract_storage_get_after_finalize() { + let contracts = (0..4).map(|i| ContractCode::new(vec![i], None)).collect_vec(); + let (existing_contracts, missing_contracts) = contracts.split_at(2); + + // Insert old contracts (already deployed contracts) into the storage. + let mut mock_storage = MockTrieStorage::new(); + for contract in existing_contracts.iter() { + mock_storage.insert(*contract.hash(), contract.code().to_vec().into()); + } + + let contract_storage = ContractStorage::new(Arc::new(mock_storage)); + + // Make the same `get` calls before and after finalizing the tracker. + // Both should return the same results for existing and missing contracts. + + for contract in existing_contracts.iter() { + assert_eq!(contract_storage.get(*contract.hash()).unwrap().hash(), contract.hash()); + } + for contract in missing_contracts.iter() { + assert!(contract_storage.get(*contract.hash()).is_none()); + } + + let contract_storage_clone = contract_storage.clone(); + let _ = contract_storage.finalize(); + + for contract in existing_contracts.iter() { + assert_eq!( + contract_storage_clone.get(*contract.hash()).unwrap().hash(), + contract.hash() + ); + } + for contract in missing_contracts.iter() { + assert!(contract_storage_clone.get(*contract.hash()).is_none()); + } + } +}