test(starknet_class_manager_types): add MemoryClassManagerClient for testing

crates/papyrus_p2p_sync/Cargo.toml

@@ -39,6 +39,7 @@ lazy_static.workspace = true
 papyrus_network = { workspace = true, features = ["testing"] }
 papyrus_protobuf = { workspace = true, features = ["testing"] }
 papyrus_storage = { workspace = true, features = ["testing"] }
+starknet_class_manager_types = { workspace = true, features = ["testing"] }
 static_assertions.workspace = true
 tokio = { workspace = true, features = ["test-util"] }
 

crates/papyrus_p2p_sync/src/client/class_test.rs

@@ -11,11 +11,12 @@ use papyrus_protobuf::sync::{
     Query,
     StateDiffChunk,
 };
-use papyrus_storage::class::ClassStorageReader;
+use papyrus_storage::class_manager::ClassManagerStorageReader;
 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,
@@ -98,19 +99,19 @@ 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(),
-                        reader.begin_ro_txn().unwrap().get_class_marker().unwrap().0
+                        reader.begin_ro_txn().unwrap().get_class_manager_block_marker().unwrap().0
                     );
                 }
                 .boxed()
             })));
             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| {
+        actions.push(Action::CheckStorage(Box::new(move |(reader, class_manager_client)| {
             async move {
                 let block_number = BlockNumber(i.try_into().unwrap());
                 wait_for_marker(
@@ -122,18 +123,18 @@ async fn class_basic_flow() {
                 )
                 .await;
 
-                let txn = reader.begin_ro_txn().unwrap();
                 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 = txn.get_class(&class_hash).unwrap().unwrap();
+                            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 =
-                                txn.get_deprecated_class(&class_hash).unwrap().unwrap();
-                            assert_eq!(actual_class, expected_class.clone());
+                                class_manager_client.get_executable(class_hash).await.unwrap();
+                            assert_eq!(actual_class, ContractClass::V0(expected_class.clone()));
                         }
                     }
                 }

crates/papyrus_p2p_sync/src/client/header_test.rs

@@ -213,7 +213,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/state_diff_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(

crates/papyrus_p2p_sync/src/client/test.rs

@@ -38,7 +38,7 @@ async fn receive_block_internally() {
         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,
@@ -113,7 +113,7 @@ async fn receive_blocks_out_of_order() {
             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,
@@ -193,7 +193,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 +212,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_utils.rs

@@ -25,7 +25,7 @@ use papyrus_protobuf::sync::{
     TransactionQuery,
 };
 use papyrus_storage::body::BodyStorageReader;
-use papyrus_storage::class::ClassStorageReader;
+use papyrus_storage::class_manager::ClassManagerStorageReader;
 use papyrus_storage::header::HeaderStorageReader;
 use papyrus_storage::state::StateStorageReader;
 use papyrus_storage::test_utils::get_test_storage;
@@ -44,7 +44,8 @@ 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::EmptyClassManagerClient;
+use starknet_class_manager_types::test_utils::MemoryClassManagerClient;
+use starknet_class_manager_types::SharedClassManagerClient;
 use starknet_state_sync_types::state_sync_types::SyncBlock;
 use starknet_types_core::felt::Felt;
 use tokio::sync::oneshot;
@@ -112,8 +113,7 @@ pub fn setup() -> TestArgs {
         transaction_sender,
         class_sender,
     };
-    // TODO(noamsp): use MockClassManagerClient instead
-    let class_manager_client = Arc::new(EmptyClassManagerClient);
+    let class_manager_client = Arc::new(MemoryClassManagerClient::new());
     let p2p_sync = P2pSyncClient::new(
         p2p_sync_config,
         storage_reader.clone(),
@@ -161,7 +161,9 @@ 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) -> BoxFuture<'static, ()>>),
+    CheckStorage(
+        Box<dyn FnOnce((StorageReader, SharedClassManagerClient)) -> BoxFuture<'static, ()>>,
+    ),
     /// Check that a report was sent on the current header query.
     ValidateReportSent(DataType),
     /// Sends an internal block to the sync.
@@ -200,15 +202,14 @@ 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);
-    // TODO(noamsp): use MockClassManagerClient instead
-    let class_manager_client = Arc::new(EmptyClassManagerClient);
+    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,
+        class_manager_client.clone(),
     );
 
     let mut headers_current_query_responses_manager = None;
@@ -274,7 +275,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()).await;
+                        check_storage_fn((storage_reader.clone(), class_manager_client.clone())).await;
                     }
                     Action::ValidateReportSent(DataType::Header) => {
                         let responses_manager = headers_current_query_responses_manager.take()
@@ -386,7 +387,7 @@ pub(crate) async fn wait_for_marker(
             DataType::Header => txn.get_header_marker().unwrap(),
             DataType::Transaction => txn.get_body_marker().unwrap(),
             DataType::StateDiff => txn.get_state_marker().unwrap(),
-            DataType::Class => txn.get_class_marker().unwrap(),
+            DataType::Class => txn.get_class_manager_block_marker().unwrap(),
         };
 
         if storage_marker >= expected_marker {

crates/papyrus_p2p_sync/src/client/transaction_test.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(

crates/papyrus_p2p_sync/src/server/test.rs

@@ -26,7 +26,7 @@ use papyrus_protobuf::sync::{
     TransactionQuery,
 };
 use papyrus_storage::body::BodyStorageWriter;
-use papyrus_storage::class::ClassStorageWriter;
+use papyrus_storage::class_manager::ClassManagerStorageWriter;
 use papyrus_storage::header::{HeaderStorageReader, HeaderStorageWriter};
 use papyrus_storage::state::StateStorageWriter;
 use papyrus_storage::test_utils::get_test_storage;
@@ -50,7 +50,8 @@ use starknet_api::transaction::{
     TransactionHash,
     TransactionOutput,
 };
-use starknet_class_manager_types::EmptyClassManagerClient;
+use starknet_class_manager_types::test_utils::MemoryClassManagerClient;
+use starknet_class_manager_types::SharedClassManagerClient;
 
 use super::{split_thin_state_diff, FetchBlockData, P2pSyncServer, P2pSyncServerChannels};
 use crate::server::register_query;
@@ -339,6 +340,13 @@ async fn run_test<T, F, TQuery>(
     // put some data in the storage.
     insert_to_storage_test_blocks_up_to(&mut storage_writer);
 
+    // put classes into class manager and update marker in storage
+    insert_to_class_manager_test_blocks_up_to(
+        &mut storage_writer,
+        p2p_sync_server.class_manager_client.clone(),
+    )
+    .await;
+
     let block_number = BlockNumber(start_block_number.try_into().unwrap());
     let start_block = match start_block_type {
         StartBlockType::Hash => BlockHashOrNumber::Hash(
@@ -364,9 +372,12 @@ async fn run_test<T, F, TQuery>(
     let (server_query_manager, _report_sender, response_reciever) =
         create_test_server_query_manager(query);
 
-    // TODO(noamsp): use MockClassManagerClient instead
-    let class_manager_client = Arc::new(EmptyClassManagerClient);
-    register_query::<T, TQuery>(storage_reader, server_query_manager, class_manager_client, "test");
+    register_query::<T, TQuery>(
+        storage_reader,
+        server_query_manager,
+        p2p_sync_server.class_manager_client.clone(),
+        "test",
+    );
 
     // run p2p_sync_server and collect query results.
     tokio::select! {
@@ -414,8 +425,7 @@ fn setup() -> TestArgs {
         event_receiver,
     };
 
-    // TODO(noamsp): use MockClassManagerClient instead
-    let class_manager_client = Arc::new(EmptyClassManagerClient);
+    let class_manager_client = Arc::new(MemoryClassManagerClient::new());
 
     let p2p_sync_server = super::P2pSyncServer::new(
         storage_reader.clone(),
@@ -433,6 +443,7 @@ fn setup() -> TestArgs {
         event_sender,
     }
 }
+
 use starknet_api::core::ClassHash;
 fn insert_to_storage_test_blocks_up_to(storage_writer: &mut StorageWriter) {
     for i in 0..NUM_OF_BLOCKS {
@@ -445,14 +456,6 @@ fn insert_to_storage_test_blocks_up_to(storage_writer: &mut StorageWriter) {
             },
             ..Default::default()
         };
-        let classes_with_hashes = CLASSES_WITH_HASHES[i]
-            .iter()
-            .map(|(class_hash, contract_class)| (*class_hash, contract_class))
-            .collect::<Vec<_>>();
-        let deprecated_classes_with_hashes = DEPRECATED_CLASSES_WITH_HASHES[i]
-            .iter()
-            .map(|(class_hash, contract_class)| (*class_hash, contract_class))
-            .collect::<Vec<_>>();
         storage_writer
             .begin_rw_txn()
             .unwrap()
@@ -467,7 +470,42 @@ fn insert_to_storage_test_blocks_up_to(storage_writer: &mut StorageWriter) {
             .append_body(block_number, BlockBody{transactions: TXS[i].clone(),
                 transaction_outputs: TX_OUTPUTS[i].clone(),
                 transaction_hashes: TX_HASHES[i].clone(),}).unwrap()
-            .append_classes(block_number, &classes_with_hashes, &deprecated_classes_with_hashes)
+            .commit()
+            .unwrap();
+    }
+}
+
+async fn insert_to_class_manager_test_blocks_up_to(
+    storage_writer: &mut StorageWriter,
+    class_manager_client: SharedClassManagerClient,
+) {
+    for i in 0..NUM_OF_BLOCKS {
+        let block_number = BlockNumber(i.try_into().unwrap());
+        let classes_with_hashes = CLASSES_WITH_HASHES[i]
+            .iter()
+            .map(|(class_hash, contract_class)| (*class_hash, contract_class))
+            .collect::<Vec<_>>();
+
+        for (class_hash, contract_class) in classes_with_hashes {
+            class_manager_client.add_class(class_hash, contract_class.clone()).await.unwrap();
+        }
+
+        let deprecated_classes_with_hashes = DEPRECATED_CLASSES_WITH_HASHES[i]
+            .iter()
+            .map(|(class_hash, contract_class)| (*class_hash, contract_class))
+            .collect::<Vec<_>>();
+
+        for (class_hash, contract_class) in deprecated_classes_with_hashes {
+            class_manager_client
+                .add_deprecated_class(class_hash, contract_class.clone())
+                .await
+                .unwrap();
+        }
+
+        storage_writer
+            .begin_rw_txn()
+            .unwrap()
+            .update_class_manager_block_marker(&block_number)
             .unwrap()
             .commit()
             .unwrap();

crates/starknet_class_manager_types/Cargo.toml

@@ -5,6 +5,9 @@ license.workspace = true
 repository.workspace = true
 version.workspace = true
 
+[features]
+testing = []
+
 [lints]
 workspace = true
 

crates/starknet_class_manager_types/src/lib.rs

@@ -1,5 +1,6 @@
+#[cfg(any(feature = "testing", test))]
+pub mod test_utils;
 pub mod transaction_converter;
-
 use std::sync::Arc;
 
 use async_trait::async_trait;