From c2be47fb499bca3d320c5b161111ad98c9b3ee71 Mon Sep 17 00:00:00 2001
From: AvivYossef-starkware <aviv.y@starkware.co>
Date: Wed, 17 Jul 2024 16:40:37 +0300
Subject: [PATCH] refactor: use mock test utils in update skeleton tests

---
 .../internal_test_utils.rs                     |  4 ++++
 .../create_tree_helper_test.rs                 | 18 +++++++++---------
 .../updated_skeleton_tree/tree_test.rs         | 12 ++++++------
 3 files changed, 19 insertions(+), 15 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
index d83054b9..c439bb71 100644
--- a/crates/committer/src/patricia_merkle_tree/internal_test_utils.rs
+++ b/crates/committer/src/patricia_merkle_tree/internal_test_utils.rs
@@ -26,6 +26,8 @@ use ethnum::U256;
 use rand::rngs::ThreadRng;
 use rstest::{fixture, rstest};
 
+use crate::patricia_merkle_tree::filled_tree::tree::FilledTreeImpl;
+
 #[derive(Debug, PartialEq, Clone, Copy, Default, Eq)]
 pub(crate) struct MockLeaf(pub(crate) Felt);
 
@@ -76,6 +78,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
index 6d0be571..8441b493 100644
--- 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;
@@ -501,9 +500,11 @@ fn test_update_node_in_nonempty_tree(
 #[case::non_empty_tree(HashOutput(Felt::from(77_u128)))]
 #[tokio::test]
 async fn test_update_non_modified_storage_tree(#[case] root_hash: HashOutput) {
+    use crate::patricia_merkle_tree::internal_test_utils::OriginalSkeletonMockTrieConfig;
+
     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 +513,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
index fefda311..f08389f2 100644
--- 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,11 @@ 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::{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 +143,17 @@ 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>) {
+    use crate::patricia_merkle_tree::internal_test_utils::OriginalSkeletonMockTrieConfig;
+
     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();