From 7ccede7e2b0c1a762f1bd6c49ba771b7e373f2be Mon Sep 17 00:00:00 2001
From: grantkee <49913008+grantkee@users.noreply.github.com>
Date: Mon, 1 May 2023 09:36:51 -0500
Subject: [PATCH] feat(mock) `ExtendedAccount` storage method (#1888)

---
 crates/rpc/rpc/src/eth/api/state.rs           | 44 +++++++++++++++++++
 .../storage/provider/src/test_utils/mock.rs   | 10 +++++
 2 files changed, 54 insertions(+)

diff --git a/crates/rpc/rpc/src/eth/api/state.rs b/crates/rpc/rpc/src/eth/api/state.rs
index f8462de4af..cb7b71472a 100644
--- a/crates/rpc/rpc/src/eth/api/state.rs
+++ b/crates/rpc/rpc/src/eth/api/state.rs
@@ -140,3 +140,47 @@ where
         Ok(proof)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::eth::cache::EthStateCache;
+    use reth_primitives::{StorageKey, StorageValue};
+    use reth_provider::test_utils::{ExtendedAccount, MockEthProvider, NoopProvider};
+    use reth_transaction_pool::test_utils::testing_pool;
+    use std::collections::HashMap;
+
+    #[tokio::test]
+    async fn test_storage() {
+        // === Noop ===
+        let pool = testing_pool();
+
+        let eth_api = EthApi::new(
+            NoopProvider::default(),
+            pool.clone(),
+            (),
+            EthStateCache::spawn(NoopProvider::default(), Default::default()),
+        );
+        let address = Address::random();
+        let storage = eth_api.storage_at(address, U256::ZERO, None).unwrap();
+        assert_eq!(storage, U256::ZERO.into());
+
+        // === Mock ===
+        let mock_provider = MockEthProvider::default();
+        let storage_value = StorageValue::from(1337);
+        let storage_key = StorageKey::random();
+        let storage = HashMap::from([(storage_key, storage_value)]);
+        let account = ExtendedAccount::new(0, U256::ZERO).extend_storage(storage);
+        mock_provider.add_account(address, account);
+
+        let eth_api = EthApi::new(
+            mock_provider.clone(),
+            pool.clone(),
+            (),
+            EthStateCache::spawn(mock_provider, Default::default()),
+        );
+
+        let storage = eth_api.storage_at(address, storage_key.into(), None).unwrap();
+        assert_eq!(storage, storage_value.into());
+    }
+}
diff --git a/crates/storage/provider/src/test_utils/mock.rs b/crates/storage/provider/src/test_utils/mock.rs
index 72a4ddb205..d8e9ed5f87 100644
--- a/crates/storage/provider/src/test_utils/mock.rs
+++ b/crates/storage/provider/src/test_utils/mock.rs
@@ -54,6 +54,16 @@ impl ExtendedAccount {
         self.bytecode = Some(Bytecode::new_raw(bytecode.into()));
         self
     }
+
+    /// Add storage to the extended account. If the storage key is already present,
+    /// the value is updated.
+    pub fn extend_storage(
+        mut self,
+        storage: impl IntoIterator<Item = (StorageKey, StorageValue)>,
+    ) -> Self {
+        self.storage.extend(storage);
+        self
+    }
 }
 
 impl MockEthProvider {