refactor: use mock test utils in update skeleton tests

starkware-libs · Jul 22, 2024 · 172dc30 · 172dc30
1 parent 2e85680
commit 172dc30
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 17 deletions.
diff --git a/crates/committer/src/patricia_merkle_tree/internal_test_utils.rs b/crates/committer/src/patricia_merkle_tree/internal_test_utils.rs
@@ -1,10 +1,10 @@
 use std::sync::Arc;
 
 use crate::felt::Felt;
+use crate::generate_trie_config;
 use crate::hash::hash_trait::HashOutput;
 use crate::patricia_merkle_tree::external_test_utils::get_random_u256;
-
-use crate::generate_trie_config;
+use crate::patricia_merkle_tree::filled_tree::tree::FilledTreeImpl;
 use crate::patricia_merkle_tree::node_data::errors::LeafResult;
 use crate::patricia_merkle_tree::node_data::inner_node::NodeData;
 use crate::patricia_merkle_tree::node_data::inner_node::{EdgePathLength, PathToBottom};
@@ -76,6 +76,8 @@ impl TreeHashFunction<MockLeaf> for TreeHashFunctionImpl {
 
 generate_trie_config!(OriginalSkeletonMockTrieConfig, MockLeaf);
 
+pub(crate) type MockTrie = FilledTreeImpl<MockLeaf>;
+
 struct MockHashFunction;
 
 impl HashFunction for MockHashFunction {

diff --git a/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/create_tree_helper_test.rs b/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/create_tree_helper_test.rs
@@ -1,7 +1,6 @@
-use crate::block_committer::input::StarknetStorageValue;
 use crate::patricia_merkle_tree::filled_tree::tree::FilledTree;
-use crate::patricia_merkle_tree::filled_tree::tree::StorageTrie;
-use crate::patricia_merkle_tree::original_skeleton_tree::config::OriginalSkeletonStorageTrieConfig;
+use crate::patricia_merkle_tree::internal_test_utils::MockLeaf;
+use crate::patricia_merkle_tree::internal_test_utils::MockTrie;
 use crate::patricia_merkle_tree::types::SortedLeafIndices;
 use ethnum::{uint, U256};
 use pretty_assertions::assert_eq;
@@ -11,6 +10,7 @@ use std::sync::Arc;
 
 use crate::felt::Felt;
 use crate::hash::hash_trait::HashOutput;
+use crate::patricia_merkle_tree::internal_test_utils::OriginalSkeletonMockTrieConfig;
 use crate::patricia_merkle_tree::internal_test_utils::{
     as_fully_indexed, get_initial_updated_skeleton, small_tree_index_to_full,
 };
@@ -502,8 +502,8 @@ fn test_update_node_in_nonempty_tree(
 #[tokio::test]
 async fn test_update_non_modified_storage_tree(#[case] root_hash: HashOutput) {
     let empty_map = HashMap::new();
-    let config = OriginalSkeletonStorageTrieConfig::new(&empty_map, false);
-    let mut original_skeleton_tree = OriginalSkeletonTreeImpl::create_impl::<StarknetStorageValue>(
+    let config = OriginalSkeletonMockTrieConfig::new(&empty_map, false);
+    let mut original_skeleton_tree = OriginalSkeletonTreeImpl::create_impl::<MockLeaf>(
         &MapStorage::default(),
         root_hash,
         SortedLeafIndices::new(&mut []),
@@ -512,9 +512,8 @@ async fn test_update_non_modified_storage_tree(#[case] root_hash: HashOutput) {
     .unwrap();
     let updated =
         UpdatedSkeletonTreeImpl::create(&mut original_skeleton_tree, &HashMap::new()).unwrap();
-    let filled =
-        StorageTrie::create::<TreeHashFunctionImpl>(Arc::new(updated), Arc::new(empty_map))
-            .await
-            .unwrap();
+    let filled = MockTrie::create::<TreeHashFunctionImpl>(Arc::new(updated), Arc::new(empty_map))
+        .await
+        .unwrap();
     assert_eq!(root_hash, filled.get_root_hash());
 }
diff --git a/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/tree_test.rs b/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/tree_test.rs
@@ -2,13 +2,12 @@ use std::collections::HashMap;
 
 use rstest::{fixture, rstest};
 
-use crate::block_committer::input::StarknetStorageValue;
 use crate::felt::Felt;
 use crate::hash::hash_trait::HashOutput;
-use crate::patricia_merkle_tree::internal_test_utils::get_initial_updated_skeleton;
+use crate::patricia_merkle_tree::internal_test_utils::OriginalSkeletonMockTrieConfig;
+use crate::patricia_merkle_tree::internal_test_utils::{get_initial_updated_skeleton, MockLeaf};
 use crate::patricia_merkle_tree::node_data::inner_node::PathToBottom;
 use crate::patricia_merkle_tree::node_data::leaf::{LeafModifications, SkeletonLeaf};
-use crate::patricia_merkle_tree::original_skeleton_tree::config::OriginalSkeletonStorageTrieConfig;
 use crate::patricia_merkle_tree::original_skeleton_tree::node::OriginalSkeletonNode;
 use crate::patricia_merkle_tree::original_skeleton_tree::tree::{
     OriginalSkeletonTree, OriginalSkeletonTreeImpl,
@@ -145,15 +144,15 @@ fn test_updated_skeleton_tree_impl_create(
 
 #[rstest]
 #[case::empty_modifications(HashMap::new())]
-#[case::non_empty_modifications(HashMap::from([(NodeIndex::FIRST_LEAF + NodeIndex::from(7), StarknetStorageValue::default())]))]
-fn test_updated_empty_tree(#[case] modifications: LeafModifications<StarknetStorageValue>) {
+#[case::non_empty_modifications(HashMap::from([(NodeIndex::FIRST_LEAF + NodeIndex::from(7), MockLeaf::default())]))]
+fn test_updated_empty_tree(#[case] modifications: LeafModifications<MockLeaf>) {
     let storage: MapStorage = HashMap::new().into();
     let mut indices: Vec<NodeIndex> = modifications.keys().copied().collect();
     let mut original_skeleton = OriginalSkeletonTreeImpl::create(
         &storage,
         HashOutput::ROOT_OF_EMPTY_TREE,
         SortedLeafIndices::new(&mut indices),
-        &OriginalSkeletonStorageTrieConfig::new(&modifications, false),
+        &OriginalSkeletonMockTrieConfig::new(&modifications, false),
     )
     .unwrap();