From 1b2c3ea459c8335393889c348a6b82a643d88525 Mon Sep 17 00:00:00 2001
From: mariari <mariari@protonmail.com>
Date: Tue, 17 Jan 2023 19:07:39 +0800
Subject: [PATCH] Add token initialization to the test storage

---
 core/src/ledger/storage/mod.rs | 52 ++++++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/core/src/ledger/storage/mod.rs b/core/src/ledger/storage/mod.rs
index 697a5df66c..199965a4e4 100644
--- a/core/src/ledger/storage/mod.rs
+++ b/core/src/ledger/storage/mod.rs
@@ -1242,10 +1242,13 @@ impl From<MerkleTreeError> for Error {
 /// Helpers for testing components that depend on storage
 #[cfg(any(test, feature = "testing"))]
 pub mod testing {
+    use rust_decimal_macros::dec;
+
     use super::mockdb::MockDB;
     use super::*;
     use crate::ledger::storage::traits::Sha256Hasher;
     use crate::types::address;
+    use crate::types::token::parameters;
     /// Storage with a mock DB for testing
     pub type TestStorage = Storage<MockDB, Sha256Hasher>;
 
@@ -1282,6 +1285,54 @@ pub mod testing {
             }
         }
     }
+
+    impl TestStorage {
+        /// Initializes all tokens in the test storage
+        pub fn initalize_tokens(
+            &mut self,
+            total_supply: token::Amount,
+            total_supply_in_masp: token::Amount,
+        ) {
+            let masp_rewards = address::masp_rewards();
+            let masp_addr = address::masp();
+            for addr in masp_rewards.keys() {
+                parameters::Parameters::init_storage(
+                    &parameters::Parameters::default(),
+                    addr,
+                    self,
+                );
+                let initial_inflation: u64 = 1;
+                let initial_locked_ratio: Decimal = dec!(0.1);
+
+                StorageWrite::write(
+                    self,
+                    &token::last_inflation(addr),
+                    initial_inflation,
+                )
+                .expect("Should not fail to put a test inflation source");
+                StorageWrite::write(
+                    self,
+                    &token::last_locked_ratio(addr),
+                    initial_locked_ratio,
+                )
+                .expect("Should not fail to put a test inflation source");
+
+                self.write(
+                    &token::total_supply_key(addr),
+                    total_supply.try_to_vec().unwrap(),
+                )
+                .expect("Should not fail to put a test total supply");
+                self.write(
+                    &token::balance_key(addr, &masp_addr),
+                    total_supply_in_masp.try_to_vec().unwrap(),
+                )
+                .expect(
+                    "Should not fail to put a test the total supply in the \
+                     MASP",
+                );
+            }
+        }
+    }
 }
 
 #[cfg(test)]
@@ -1379,6 +1430,7 @@ mod tests {
             };
             parameters.init_storage(&mut storage);
 
+            storage.initalize_tokens(token::Amount::from(1000), token::Amount::from(500));
             let epoch_before = storage.last_epoch;
             assert_eq!(epoch_before, storage.block.epoch);