feat : added tests for refactor_state_update function and refactors

crates/orchestrator/src/jobs/da_job/mod.rs

@@ -396,33 +396,33 @@ fn da_word(class_flag: bool, nonce_change: Option<Felt>, num_changes: u64) -> Fe
 }
 
 fn refactor_state_update(state_update: &mut StateDiff) {
-    let addresses_in_nonces: Vec<Felt> = state_update.nonces.clone().iter().map(|item| item.contract_address).collect();
-    let addresses_in_storage_diffs: Vec<Felt> =
-        state_update.storage_diffs.clone().iter().map(|item| item.address).collect();
+    // Create set of all addresses that need storage diffs
+    let mut required_addresses: HashSet<Felt> = HashSet::new();
 
-    let address_to_insert = find_unique_addresses(addresses_in_nonces, addresses_in_storage_diffs.clone());
+    // First add all existing storage addresses
+    required_addresses.extend(state_update.storage_diffs.iter().map(|item| item.address));
 
-    for address in address_to_insert {
-        state_update.storage_diffs.push(ContractStorageDiffItem { address, storage_entries: vec![] })
-    }
-
-    let addresses_in_deployed_contracts =
-        state_update.deployed_contracts.clone().iter().map(|item| item.address).collect();
-    let addresses_to_insert = find_unique_addresses(addresses_in_deployed_contracts, addresses_in_storage_diffs);
+    // Add addresses from nonces
+    required_addresses.extend(state_update.nonces.iter().map(|item| item.contract_address));
 
-    for address in addresses_to_insert {
-        state_update.storage_diffs.push(ContractStorageDiffItem { address, storage_entries: vec![] })
-    }
-}
+    // Add addresses from deployed contracts
+    required_addresses.extend(state_update.deployed_contracts.iter().map(|item| item.address));
 
-fn find_unique_addresses(nonce_addresses: Vec<Felt>, storage_diff_addresses: Vec<Felt>) -> Vec<Felt> {
-    let storage_set: HashSet<_> = storage_diff_addresses.into_iter().collect();
+    // Create a set of existing storage addresses for quick lookup
+    let existing_storage_addresses: HashSet<Felt> =
+        state_update.storage_diffs.iter().map(|item| item.address).collect();
 
-    nonce_addresses.into_iter().filter(|addr| !storage_set.contains(addr)).collect()
+    // Add storage diffs only for addresses that don't already have them
+    for address in required_addresses {
+        if !existing_storage_addresses.contains(&address) {
+            state_update.storage_diffs.push(ContractStorageDiffItem { address, storage_entries: Vec::new() });
+        }
+    }
 }
 
 #[cfg(test)]
 pub mod test {
+    use std::collections::HashSet;
     use std::fs;
     use std::fs::File;
     use std::io::Read;
@@ -435,12 +435,14 @@ pub mod test {
     use majin_blob_types::serde;
     use rstest::rstest;
     use serde_json::json;
-    use starknet::core::types::{Felt, StateUpdate};
+    use starknet::core::types::{
+        ContractStorageDiffItem, DeployedContractItem, Felt, NonceUpdate, StateDiff, StateUpdate, StorageEntry,
+    };
     use starknet::providers::jsonrpc::HttpTransport;
     use starknet::providers::JsonRpcClient;
     use url::Url;
 
-    use crate::jobs::da_job::da_word;
+    use crate::jobs::da_job::{da_word, refactor_state_update};
 
     /// Tests `da_word` function with various inputs for class flag, new nonce, and number of
     /// changes. Verifies that `da_word` produces the correct Felt based on the provided
@@ -494,8 +496,12 @@ pub mod test {
         "src/tests/jobs/da_job/test_data/test_blob/671070.txt",
         "src/tests/jobs/da_job/test_data/nonces/671070.txt"
     )]
+    // Block from pragma madara and orch test run. Here we faced an issue where our
+    // blob building logic was not able to take the contract addresses from
+    // `deployed_contracts` field in state diff from state update. This test case
+    // was added after the fix
     #[case(
-        178, // Block from pragma madara and orch test run
+        178,
         "src/tests/jobs/da_job/test_data/state_update/178.txt",
         "src/tests/jobs/da_job/test_data/test_blob/178.txt",
         "src/tests/jobs/da_job/test_data/nonces/178.txt"
@@ -584,6 +590,71 @@ pub mod test {
         assert_eq!(data, deserialize_data);
     }
 
+    #[rstest]
+    #[case(vec![], vec![], vec![], 0)] // Empty case
+    #[case(
+        vec![(Felt::from(1), Felt::from(10)), (Felt::from(2), Felt::from(20))],
+        vec![],
+        vec![],
+        2
+    )] // Only nonces
+    #[case(
+        vec![],
+        vec![],
+        vec![(Felt::from(1), vec![1]), (Felt::from(2), vec![2])],
+        2
+    )] // Only deployed
+    #[case(
+        vec![(Felt::from(1), Felt::from(10))],
+        vec![(Felt::from(1), vec![(Felt::from(1), Felt::from(100))])],
+        vec![(Felt::from(1), vec![1])],
+        1
+    )] // Overlapping addresses
+    #[case(
+        vec![(Felt::from(1), Felt::from(10)), (Felt::from(1), Felt::from(20))],
+        vec![],
+        vec![(Felt::from(1), vec![1]), (Felt::from(1), vec![2])],
+        1
+    )] // Duplicate addresses
+    fn test_refactor_state_update(
+        #[case] nonces: Vec<(Felt, Felt)>,
+        #[case] storage_diffs: Vec<(Felt, Vec<(Felt, Felt)>)>,
+        #[case] deployed_contracts: Vec<(Felt, Vec<u8>)>,
+        #[case] expected_storage_count: usize,
+    ) {
+        // Create initial state
+        let mut state_diff = create_state_diff(nonces, storage_diffs.clone(), deployed_contracts);
+        let initial_storage = state_diff.storage_diffs.clone();
+
+        // Run the function
+        refactor_state_update(&mut state_diff);
+
+        // Verify results
+        assert!(verify_addresses_have_storage_diffs(&state_diff));
+
+        // Verify no duplicate addresses in storage_diffs
+        let unique_addresses: HashSet<_> = state_diff.storage_diffs.iter().map(|item| &item.address).collect();
+        assert_eq!(
+            unique_addresses.len(),
+            state_diff.storage_diffs.len(),
+            "Number of unique addresses should match number of storage diffs"
+        );
+        assert_eq!(
+            unique_addresses.len(),
+            expected_storage_count,
+            "Number of storage diffs should match expected count"
+        );
+
+        // Verify storage preservation for existing entries
+        for orig_storage in initial_storage {
+            if let Some(current_storage) =
+                state_diff.storage_diffs.iter().find(|item| item.address == orig_storage.address)
+            {
+                assert_eq!(orig_storage.storage_entries, current_storage.storage_entries);
+            }
+        }
+    }
+
     pub(crate) fn read_state_update_from_file(file_path: &str) -> Result<StateUpdate> {
         // let file_path = format!("state_update_block_no_{}.txt", block_no);
         let mut file = File::open(file_path)?;
@@ -630,4 +701,35 @@ pub mod test {
         new_hex_chars = new_hex_chars.trim_start_matches('0').to_string();
         if new_hex_chars.is_empty() { "0x0".to_string() } else { format!("0x{}", new_hex_chars) }
     }
+
+    fn create_state_diff(
+        nonces: Vec<(Felt, Felt)>,
+        storage_diffs: Vec<(Felt, Vec<(Felt, Felt)>)>,
+        deployed_contracts: Vec<(Felt, Vec<u8>)>,
+    ) -> StateDiff {
+        StateDiff {
+            nonces: nonces.into_iter().map(|(addr, nonce)| NonceUpdate { contract_address: addr, nonce }).collect(),
+            storage_diffs: storage_diffs
+                .into_iter()
+                .map(|(addr, entries)| ContractStorageDiffItem {
+                    address: addr,
+                    storage_entries: entries.into_iter().map(|(key, value)| StorageEntry { key, value }).collect(),
+                })
+                .collect(),
+            deprecated_declared_classes: vec![],
+            declared_classes: vec![],
+            deployed_contracts: deployed_contracts
+                .into_iter()
+                .map(|(addr, _code)| DeployedContractItem { address: addr, class_hash: Default::default() })
+                .collect(),
+            replaced_classes: vec![],
+        }
+    }
+
+    fn verify_addresses_have_storage_diffs(state_diff: &StateDiff) -> bool {
+        let storage_addresses: HashSet<_> = state_diff.storage_diffs.iter().map(|item| &item.address).collect();
+
+        state_diff.nonces.iter().all(|item| storage_addresses.contains(&item.contract_address))
+            && state_diff.deployed_contracts.iter().all(|item| storage_addresses.contains(&item.address))
+    }
 }