From b0f2ac346452e61709c4d0699f2f9d403381e13d Mon Sep 17 00:00:00 2001
From: eitanm-starkware <eitan.m@starkware.co>
Date: Mon, 17 Jun 2024 14:28:04 +0300
Subject: [PATCH] test(protobuf): add auto impl for statediffchunk

---
 Cargo.lock                                    |  2 ++
 crates/papyrus_common/src/state.rs            |  1 +
 crates/papyrus_protobuf/Cargo.toml            |  8 +++++
 .../src/converters/state_diff.rs              |  1 +
 crates/papyrus_protobuf/src/sync.rs           | 36 +++++++++++++++++--
 5 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 9b96c88ea2..0a3a3ee28b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -6167,6 +6167,8 @@ dependencies = [
  "prost",
  "prost-build",
  "prost-types",
+ "rand 0.8.5",
+ "rand_chacha 0.3.1",
  "starknet-types-core",
  "starknet_api",
  "test_utils",
diff --git a/crates/papyrus_common/src/state.rs b/crates/papyrus_common/src/state.rs
index 4254f8133d..badf6d7ba0 100644
--- a/crates/papyrus_common/src/state.rs
+++ b/crates/papyrus_common/src/state.rs
@@ -32,6 +32,7 @@ pub struct ReplacedClass {
     pub class_hash: ClassHash,
 }
 
+// TODO: move to used crate
 pub fn create_random_state_diff(rng: &mut impl RngCore) -> ThinStateDiff {
     let contract0 = ContractAddress::from(rng.next_u64());
     let contract1 = ContractAddress::from(rng.next_u64());
diff --git a/crates/papyrus_protobuf/Cargo.toml b/crates/papyrus_protobuf/Cargo.toml
index 1a0e7e6190..848f304954 100644
--- a/crates/papyrus_protobuf/Cargo.toml
+++ b/crates/papyrus_protobuf/Cargo.toml
@@ -5,6 +5,9 @@ edition.workspace = true
 repository.workspace = true
 license-file.workspace = true
 
+[features]
+testing = ["rand", "rand_chacha", "test_utils"]
+
 [dependencies]
 bytes.workspace = true
 indexmap.workspace = true
@@ -12,12 +15,17 @@ lazy_static.workspace = true
 primitive-types.workspace = true
 prost.workspace = true
 prost-types.workspace = true
+rand = { workspace = true, optional = true }
+rand_chacha = { workspace = true, optional = true }
 starknet_api.workspace = true
 starknet-types-core.workspace = true
+test_utils = { path = "../test_utils", optional = true }
 thiserror.workspace = true
 papyrus_common = { path = "../papyrus_common", version = "0.4.0-dev.2" }
 
 [dev-dependencies]
+rand.workspace = true
+rand_chacha.workspace = true
 test_utils = { path = "../test_utils" }
 
 [build-dependencies]
diff --git a/crates/papyrus_protobuf/src/converters/state_diff.rs b/crates/papyrus_protobuf/src/converters/state_diff.rs
index 8cef22838b..7ce2d042c6 100644
--- a/crates/papyrus_protobuf/src/converters/state_diff.rs
+++ b/crates/papyrus_protobuf/src/converters/state_diff.rs
@@ -39,6 +39,7 @@ impl TryFrom<protobuf::StateDiffsResponse> for DataOrFin<ThinStateDiff> {
     }
 }
 auto_impl_try_from_vec_u8!(DataOrFin<ThinStateDiff>, protobuf::StateDiffsResponse);
+auto_impl_try_from_vec_u8!(DataOrFin<StateDiffChunk>, protobuf::StateDiffsResponse);
 
 impl TryFrom<protobuf::StateDiffsResponse> for DataOrFin<StateDiffChunk> {
     type Error = ProtobufConversionError;
diff --git a/crates/papyrus_protobuf/src/sync.rs b/crates/papyrus_protobuf/src/sync.rs
index 06d228d1ca..a62d04e296 100644
--- a/crates/papyrus_protobuf/src/sync.rs
+++ b/crates/papyrus_protobuf/src/sync.rs
@@ -5,6 +5,8 @@ use starknet_api::block::{BlockHash, BlockHeader, BlockNumber, BlockSignature};
 use starknet_api::core::{ClassHash, CompiledClassHash, ContractAddress, Nonce};
 use starknet_api::state::StorageKey;
 use starknet_types_core::felt::Felt;
+#[cfg(any(feature = "testing", test))]
+use test_utils::{auto_impl_get_test_instance, get_number_of_variants, GetTestInstance};
 
 #[derive(Debug, PartialEq, Eq, Clone, Copy, Default, Hash)]
 pub enum Direction {
@@ -49,7 +51,7 @@ pub struct SignedBlockHeader {
     pub signatures: Vec<BlockSignature>,
 }
 
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Default, Clone, PartialEq, Eq)]
 pub struct ContractDiff {
     pub contract_address: ContractAddress,
     // Has value only if the contract was deployed or replaced in this block.
@@ -59,13 +61,13 @@ pub struct ContractDiff {
     pub storage_diffs: IndexMap<StorageKey, Felt>,
 }
 
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Default, Clone, PartialEq, Eq)]
 pub struct DeclaredClass {
     pub class_hash: ClassHash,
     pub compiled_class_hash: CompiledClassHash,
 }
 
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Default, Clone, PartialEq, Eq)]
 pub struct DeprecatedDeclaredClass {
     pub class_hash: ClassHash,
 }
@@ -76,3 +78,31 @@ pub enum StateDiffChunk {
     DeclaredClass(DeclaredClass),
     DeprecatedDeclaredClass(DeprecatedDeclaredClass),
 }
+
+impl Default for StateDiffChunk {
+    fn default() -> Self {
+        Self::ContractDiff(ContractDiff::default())
+    }
+}
+
+#[cfg(any(feature = "testing", test))]
+auto_impl_get_test_instance! {
+    pub enum StateDiffChunk{
+        ContractDiff(ContractDiff) = 0,
+        DeclaredClass(DeclaredClass) = 1,
+        DeprecatedDeclaredClass(DeprecatedDeclaredClass) = 2,
+    }
+    pub struct ContractDiff{
+        pub contract_address: ContractAddress,
+        pub class_hash: Option<ClassHash>,
+        pub nonce: Option<Nonce>,
+        pub storage_diffs: IndexMap<StorageKey, Felt>,
+    }
+    pub struct DeclaredClass {
+        pub class_hash: ClassHash,
+        pub compiled_class_hash: CompiledClassHash,
+    }
+    pub struct DeprecatedDeclaredClass {
+        pub class_hash: ClassHash,
+    }
+}