From 9c7e61a90c6a0c47f88e69d643a285650c3e1571 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Thu, 21 Dec 2023 11:07:09 +0100
Subject: [PATCH] Add missing current era info setting

---
 .../src/benchmarking.rs                       |  1 -
 pallets/dapp-staking-v3/src/lib.rs            | 43 ++++++++++++++++---
 2 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/pallets/dapp-staking-migration/src/benchmarking.rs b/pallets/dapp-staking-migration/src/benchmarking.rs
index ffa9a13198..047e272d1b 100644
--- a/pallets/dapp-staking-migration/src/benchmarking.rs
+++ b/pallets/dapp-staking-migration/src/benchmarking.rs
@@ -105,7 +105,6 @@ mod benchmarks {
         }
     }
 
-    
     #[benchmark]
     fn cleanup_old_storage_success(x: Linear<1, 5>) {
         initial_config::<T>();
diff --git a/pallets/dapp-staking-v3/src/lib.rs b/pallets/dapp-staking-v3/src/lib.rs
index f9a10da274..c3a6c4021a 100644
--- a/pallets/dapp-staking-v3/src/lib.rs
+++ b/pallets/dapp-staking-v3/src/lib.rs
@@ -1974,28 +1974,48 @@ impl<
         // 0. Unwrap arguments
         let (init_era, tier_params, init_tier_config) = G::get();
 
-        // 1. Prepare active protocol state
+        // 1. Prepare init active protocol state
         let now = frame_system::Pallet::<T>::block_number();
         let voting_period_length = Pallet::<T>::blocks_per_voting_period();
 
+        let period_number = 1;
         let protocol_state = ProtocolState {
             era: init_era,
             next_era_start: now.saturating_add(voting_period_length),
             period_info: PeriodInfo {
-                number: 1,
+                number: period_number,
                 subperiod: Subperiod::Voting,
                 next_subperiod_start_era: init_era.saturating_add(1),
             },
             maintenance: true,
         };
 
-        // 2. Write necessary items into storage
+        // 2. Prepare init current era info - need to set correct eras
+        let init_era_info = EraInfo {
+            total_locked: 0,
+            unlocking: 0,
+            current_stake_amount: StakeAmount {
+                voting: 0,
+                build_and_earn: 0,
+                era: init_era,
+                period: period_number,
+            },
+            next_stake_amount: StakeAmount {
+                voting: 0,
+                build_and_earn: 0,
+                era: init_era.saturating_add(1),
+                period: period_number,
+            },
+        };
+
+        // 3. Write necessary items into storage
         ActiveProtocolState::<T>::put(protocol_state);
         StaticTierParams::<T>::put(tier_params);
         TierConfig::<T>::put(init_tier_config);
         STORAGE_VERSION.put::<Pallet<T>>();
+        CurrentEraInfo::<T>::put(init_era_info);
 
-        // 3. Emit events to make indexers happy
+        // 4. Emit events to make indexers happy
         Pallet::<T>::deposit_event(Event::<T>::NewEra { era: init_era });
         Pallet::<T>::deposit_event(Event::<T>::NewSubperiod {
             subperiod: Subperiod::Voting,
@@ -2004,13 +2024,14 @@ impl<
 
         log::info!("dApp Staking v3 storage initialized.");
 
-        T::DbWeight::get().reads_writes(1, 4)
+        T::DbWeight::get().reads_writes(2, 5)
     }
 
     #[cfg(feature = "try-runtime")]
     fn post_upgrade(_state: Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
         assert_eq!(Pallet::<T>::on_chain_storage_version(), STORAGE_VERSION);
-        assert!(ActiveProtocolState::<T>::get().maintenance);
+        let protocol_state = ActiveProtocolState::<T>::get();
+        assert!(protocol_state.maintenance);
 
         let number_of_tiers = T::NumberOfTiers::get();
 
@@ -2023,6 +2044,16 @@ impl<
         assert_eq!(tier_config.slots_per_tier.len(), number_of_tiers as usize);
         assert_eq!(tier_config.tier_thresholds.len(), number_of_tiers as usize);
 
+        let current_era_info = CurrentEraInfo::<T>::get();
+        assert_eq!(
+            current_era_info.current_stake_amount.era,
+            protocol_state.era
+        );
+        assert_eq!(
+            current_era_info.next_stake_amount.era,
+            protocol_state.era + 1
+        );
+
         Ok(())
     }
 }