revert(starknet_class_manager_types, papyrus_p2p_sync): delete MemoryClassManagerClient (#3560)

Author: ShahakShama

Cargo.lock

@@ -36,6 +36,7 @@ tracing.workspace = true
 [dev-dependencies]
 assert_matches.workspace = true
 lazy_static.workspace = true
+mockall.workspace = true
 papyrus_network = { workspace = true, features = ["testing"] }
 papyrus_protobuf = { workspace = true, features = ["testing"] }
 papyrus_storage = { workspace = true, features = ["testing"] }

crates/papyrus_p2p_sync/Cargo.toml

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 
 use futures::FutureExt;
+use mockall::predicate::eq;
 use papyrus_common::pending_classes::ApiContractClass;
 use papyrus_protobuf::sync::{
     BlockHashOrNumber,
@@ -16,14 +17,14 @@ use papyrus_test_utils::{get_rng, GetTestInstance};
 use rand::{Rng, RngCore};
 use rand_chacha::ChaCha8Rng;
 use starknet_api::block::BlockNumber;
-use starknet_api::contract_class::ContractClass;
 use starknet_api::core::{ClassHash, CompiledClassHash, EntryPointSelector};
 use starknet_api::deprecated_contract_class::{
     ContractClass as DeprecatedContractClass,
     EntryPointOffset,
     EntryPointV0,
 };
 use starknet_api::state::SierraContractClass;
+use starknet_class_manager_types::MockClassManagerClient;
 
 use super::test_utils::{
     random_header,
@@ -39,6 +40,8 @@ use super::test_utils::{
 async fn class_basic_flow() {
     let mut rng = get_rng();
 
+    let mut class_manager_client = MockClassManagerClient::new();
+
     let state_diffs_and_classes_of_blocks = [
         vec![
             create_random_state_diff_chunk_with_class(&mut rng),
@@ -51,6 +54,30 @@ async fn class_basic_flow() {
         ],
     ];
 
+    // Fill class manager client with expectations.
+    for state_diffs_and_classes in &state_diffs_and_classes_of_blocks {
+        for (state_diff, class) in state_diffs_and_classes {
+            let class_hash = state_diff.get_class_hash();
+            match class {
+                ApiContractClass::ContractClass(class) => {
+                    let compiled_class_hash = state_diff.get_compiled_class_hash();
+                    class_manager_client
+                        .expect_add_class()
+                        .times(1)
+                        .with(eq(class_hash), eq(class.clone()))
+                        .return_once(move |_, _| Ok(compiled_class_hash));
+                }
+                ApiContractClass::DeprecatedContractClass(class) => {
+                    class_manager_client
+                        .expect_add_deprecated_class()
+                        .times(1)
+                        .with(eq(class_hash), eq(class.clone()))
+                        .return_once(|_, _| Ok(()));
+                }
+            }
+        }
+    }
+
     let mut actions = vec![
         Action::RunP2pSync,
         // We already validate the header query content in other tests.
@@ -99,7 +126,7 @@ async fn class_basic_flow() {
             let class_hash = state_diff.get_class_hash();
 
             // Check that before the last class was sent, the classes aren't written.
-            actions.push(Action::CheckStorage(Box::new(move |(reader, _)| {
+            actions.push(Action::CheckStorage(Box::new(move |reader| {
                 async move {
                     assert_eq!(
                         u64::try_from(i).unwrap(),
@@ -111,7 +138,7 @@ async fn class_basic_flow() {
             actions.push(Action::SendClass(DataOrFin(Some((class.clone(), class_hash)))));
         }
         // Check that a block's classes are written before the entire query finished.
-        actions.push(Action::CheckStorage(Box::new(move |(reader, class_manager_client)| {
+        actions.push(Action::CheckStorage(Box::new(move |reader| {
             async move {
                 let block_number = BlockNumber(i.try_into().unwrap());
                 wait_for_marker(
@@ -122,22 +149,6 @@ async fn class_basic_flow() {
                     TIMEOUT_FOR_TEST,
                 )
                 .await;
-
-                for (state_diff, expected_class) in state_diffs_and_classes {
-                    let class_hash = state_diff.get_class_hash();
-                    match expected_class {
-                        ApiContractClass::ContractClass(expected_class) => {
-                            let actual_class =
-                                class_manager_client.get_sierra(class_hash).await.unwrap();
-                            assert_eq!(actual_class, expected_class.clone());
-                        }
-                        ApiContractClass::DeprecatedContractClass(expected_class) => {
-                            let actual_class =
-                                class_manager_client.get_executable(class_hash).await.unwrap();
-                            assert_eq!(actual_class, ContractClass::V0(expected_class.clone()));
-                        }
-                    }
-                }
             }
             .boxed()
         })));
@@ -149,6 +160,7 @@ async fn class_basic_flow() {
             (DataType::StateDiff, len.try_into().unwrap()),
             (DataType::Class, len.try_into().unwrap()),
         ]),
+        Some(class_manager_client),
         actions,
     )
     .await;
@@ -158,6 +170,7 @@ async fn class_basic_flow() {
 // we need to define this trait because StateDiffChunk is defined in an other crate.
 trait GetClassHash {
     fn get_class_hash(&self) -> ClassHash;
+    fn get_compiled_class_hash(&self) -> CompiledClassHash;
 }
 
 impl GetClassHash for StateDiffChunk {
@@ -170,6 +183,13 @@ impl GetClassHash for StateDiffChunk {
             _ => unreachable!(),
         }
     }
+
+    fn get_compiled_class_hash(&self) -> CompiledClassHash {
+        match self {
+            StateDiffChunk::DeclaredClass(declared_class) => declared_class.compiled_class_hash,
+            _ => unreachable!(),
+        }
+    }
 }
 
 fn create_random_state_diff_chunk_with_class(
@@ -325,6 +345,7 @@ async fn validate_class_sync_fails(
             (DataType::StateDiff, header_state_diff_lengths.len().try_into().unwrap()),
             (DataType::Class, header_state_diff_lengths.len().try_into().unwrap()),
         ]),
+        None,
         actions,
     )
     .await;

crates/papyrus_p2p_sync/src/client/class_test.rs

@@ -202,6 +202,7 @@ async fn sync_sends_new_header_query_if_it_got_partial_responses() {
 async fn wrong_block_number() {
     run_test(
         HashMap::from([(DataType::Header, 1)]),
+        None,
         vec![
             Action::RunP2pSync,
             // We already validate the query content in other tests.
@@ -213,7 +214,7 @@ async fn wrong_block_number() {
                 None,
             )))),
             Action::ValidateReportSent(DataType::Header),
-            Action::CheckStorage(Box::new(|(reader, _)| {
+            Action::CheckStorage(Box::new(|reader| {
                 async move {
                     assert_eq!(0, reader.begin_ro_txn().unwrap().get_header_marker().unwrap().0);
                 }

crates/papyrus_p2p_sync/src/client/header_test.rs

@@ -136,7 +136,7 @@ async fn state_diff_basic_flow() {
     {
         for state_diff_chunk in state_diff_chunks {
             // Check that before the last chunk was sent, the state diff isn't written.
-            actions.push(Action::CheckStorage(Box::new(move |(reader, _)| {
+            actions.push(Action::CheckStorage(Box::new(move |reader| {
                 async move {
                     assert_eq!(
                         u64::try_from(i).unwrap(),
@@ -149,7 +149,7 @@ async fn state_diff_basic_flow() {
             actions.push(Action::SendStateDiff(DataOrFin(Some(state_diff_chunk))));
         }
         // Check that a block's state diff is written before the entire query finished.
-        actions.push(Action::CheckStorage(Box::new(move |(reader, _)| {
+        actions.push(Action::CheckStorage(Box::new(move |reader| {
             async move {
                 let block_number = BlockNumber(i.try_into().unwrap());
                 wait_for_marker(
@@ -175,6 +175,7 @@ async fn state_diff_basic_flow() {
             (DataType::Header, state_diffs_and_chunks.len().try_into().unwrap()),
             (DataType::StateDiff, state_diffs_and_chunks.len().try_into().unwrap()),
         ]),
+        None,
         actions,
     )
     .await;
@@ -340,6 +341,7 @@ async fn validate_state_diff_fails(
             (DataType::Header, header_state_diff_lengths.len().try_into().unwrap()),
             (DataType::StateDiff, header_state_diff_lengths.len().try_into().unwrap()),
         ]),
+        None,
         actions,
     )
     .await;

crates/papyrus_p2p_sync/src/client/state_diff_test.rs

@@ -35,10 +35,11 @@ async fn receive_block_internally() {
 
     run_test(
         HashMap::new(),
+        None,
         vec![
             Action::SendInternalBlock(sync_block),
             Action::RunP2pSync,
-            Action::CheckStorage(Box::new(move |(reader, _)| {
+            Action::CheckStorage(Box::new(move |reader| {
                 async move {
                     wait_for_marker(
                         DataType::StateDiff,
@@ -109,11 +110,12 @@ async fn receive_blocks_out_of_order() {
 
     run_test(
         HashMap::new(),
+        None,
         vec![
             Action::SendInternalBlock(sync_block_1),
             Action::SendInternalBlock(sync_block_0),
             Action::RunP2pSync,
-            Action::CheckStorage(Box::new(move |(reader, _)| {
+            Action::CheckStorage(Box::new(move |reader| {
                 async move {
                     wait_for_marker(
                         DataType::StateDiff,
@@ -183,6 +185,7 @@ async fn receive_blocks_first_externally_and_then_internally() {
     let sync_block_1 = create_random_sync_block(BlockNumber(1), 1, rng);
     run_test(
         HashMap::from([(DataType::Header, 2)]),
+        None,
         vec![
             Action::RunP2pSync,
             // We already validate the query content in other tests.
@@ -193,7 +196,7 @@ async fn receive_blocks_first_externally_and_then_internally() {
                 None,
                 None,
             )))),
-            Action::CheckStorage(Box::new(|(reader, _)| {
+            Action::CheckStorage(Box::new(|reader| {
                 async move {
                     wait_for_marker(
                         DataType::Header,
@@ -212,7 +215,7 @@ async fn receive_blocks_first_externally_and_then_internally() {
             })),
             Action::SendInternalBlock(sync_block_0),
             Action::SendInternalBlock(sync_block_1),
-            Action::CheckStorage(Box::new(|(reader, _)| {
+            Action::CheckStorage(Box::new(|reader| {
                 async move {
                     wait_for_marker(
                         DataType::Header,

crates/papyrus_p2p_sync/src/client/test.rs

@@ -44,8 +44,7 @@ use starknet_api::core::ClassHash;
 use starknet_api::crypto::utils::Signature;
 use starknet_api::hash::StarkHash;
 use starknet_api::transaction::FullTransaction;
-use starknet_class_manager_types::test_utils::MemoryClassManagerClient;
-use starknet_class_manager_types::SharedClassManagerClient;
+use starknet_class_manager_types::MockClassManagerClient;
 use starknet_state_sync_types::state_sync_types::SyncBlock;
 use starknet_types_core::felt::Felt;
 use tokio::sync::oneshot;
@@ -113,7 +112,7 @@ pub fn setup() -> TestArgs {
         transaction_sender,
         class_sender,
     };
-    let class_manager_client = Arc::new(MemoryClassManagerClient::new());
+    let class_manager_client = Arc::new(MockClassManagerClient::new());
     let p2p_sync = P2pSyncClient::new(
         p2p_sync_config,
         storage_reader.clone(),
@@ -161,9 +160,7 @@ pub enum Action {
     SendClass(DataOrFin<(ApiContractClass, ClassHash)>),
     /// Perform custom validations on the storage. Returns back the storage reader it received as
     /// input
-    CheckStorage(
-        Box<dyn FnOnce((StorageReader, SharedClassManagerClient)) -> BoxFuture<'static, ()>>,
-    ),
+    CheckStorage(Box<dyn FnOnce(StorageReader) -> BoxFuture<'static, ()>>),
     /// Check that a report was sent on the current header query.
     ValidateReportSent(DataType),
     /// Sends an internal block to the sync.
@@ -172,7 +169,11 @@ pub enum Action {
 }
 
 // TODO(shahak): add support for state diffs, transactions and classes.
-pub async fn run_test(max_query_lengths: HashMap<DataType, u64>, actions: Vec<Action>) {
+pub async fn run_test(
+    max_query_lengths: HashMap<DataType, u64>,
+    class_manager_client: Option<MockClassManagerClient>,
+    actions: Vec<Action>,
+) {
     let p2p_sync_config = P2pSyncClientConfig {
         num_headers_per_query: max_query_lengths.get(&DataType::Header).cloned().unwrap_or(1),
         num_block_state_diffs_per_query: max_query_lengths
@@ -187,6 +188,8 @@ pub async fn run_test(max_query_lengths: HashMap<DataType, u64>, actions: Vec<Ac
         wait_period_for_new_data: WAIT_PERIOD_FOR_NEW_DATA,
         buffer_size: BUFFER_SIZE,
     };
+    let class_manager_client = class_manager_client.unwrap_or_default();
+    let class_manager_client = Arc::new(class_manager_client);
     let buffer_size = p2p_sync_config.buffer_size;
     let ((storage_reader, storage_writer), _temp_dir) = get_test_storage();
     let (header_sender, mut mock_header_network) = mock_register_sqmr_protocol_client(buffer_size);
@@ -202,14 +205,13 @@ pub async fn run_test(max_query_lengths: HashMap<DataType, u64>, actions: Vec<Ac
         class_sender,
     };
     let (mut internal_block_sender, internal_block_receiver) = mpsc::channel(buffer_size);
-    let class_manager_client = Arc::new(MemoryClassManagerClient::new());
     let p2p_sync = P2pSyncClient::new(
         p2p_sync_config,
         storage_reader.clone(),
         storage_writer,
         p2p_sync_channels,
         internal_block_receiver.boxed(),
-        class_manager_client.clone(),
+        class_manager_client,
     );
 
     let mut headers_current_query_responses_manager = None;
@@ -275,7 +277,7 @@ pub async fn run_test(max_query_lengths: HashMap<DataType, u64>, actions: Vec<Ac
                     }
                     Action::CheckStorage(check_storage_fn) => {
                         // We tried avoiding the clone here but it causes lifetime issues.
-                        check_storage_fn((storage_reader.clone(), class_manager_client.clone())).await;
+                        check_storage_fn(storage_reader.clone()).await;
                     }
                     Action::ValidateReportSent(DataType::Header) => {
                         let responses_manager = headers_current_query_responses_manager.take()

crates/papyrus_p2p_sync/src/client/test_utils.rs

@@ -85,7 +85,7 @@ async fn transaction_basic_flow() {
             .zip(transaction_outputs.iter().cloned().zip(transaction_hashes.iter().cloned()))
         {
             // Check that before the last transaction was sent, the transactions aren't written.
-            actions.push(Action::CheckStorage(Box::new(move |(reader, _)| {
+            actions.push(Action::CheckStorage(Box::new(move |reader| {
                 async move {
                     assert_eq!(i, reader.begin_ro_txn().unwrap().get_body_marker().unwrap().0);
                 }
@@ -100,7 +100,7 @@ async fn transaction_basic_flow() {
         }
 
         // Check that a block's transactions are written before the entire query finished.
-        actions.push(Action::CheckStorage(Box::new(move |(reader, _)| {
+        actions.push(Action::CheckStorage(Box::new(move |reader| {
             async move {
                 let block_number = BlockNumber(i);
                 wait_for_marker(
@@ -140,6 +140,7 @@ async fn transaction_basic_flow() {
             (DataType::Header, NUM_BLOCKS),
             (DataType::Transaction, TRANSACTION_QUERY_LENGTH),
         ]),
+        None,
         actions,
     )
     .await;