Fix staking try runtime (attempt #2) (#1056)

* Add missing imports

* Make staking compile for try-runtime tests (#1014)

* fix

* fix lack of From impl compilation error

Co-authored-by: Amar Singh <[email protected]>

Author: notlesh

pallets/parachain-staking/src/lib.rs

@@ -208,6 +208,63 @@ pub mod pallet {
 		pub state: CollatorStatus,
 	}
 
+	// Temporary manual implementation for migration testing purposes
+	impl<A: PartialEq, B: PartialEq> PartialEq for CollatorCandidate<A, B> {
+		fn eq(&self, other: &Self) -> bool {
+			let must_be_true = self.id == other.id
+				&& self.bond == other.bond
+				&& self.total_counted == other.total_counted
+				&& self.total_backing == other.total_backing
+				&& self.request == other.request
+				&& self.state == other.state;
+			if !must_be_true {
+				return false;
+			}
+			for (x, y) in self.delegators.0.iter().zip(other.delegators.0.iter()) {
+				if x != y {
+					return false;
+				}
+			}
+			for (
+				Bond {
+					owner: o1,
+					amount: a1,
+				},
+				Bond {
+					owner: o2,
+					amount: a2,
+				},
+			) in self
+				.top_delegations
+				.iter()
+				.zip(other.top_delegations.iter())
+			{
+				if o1 != o2 || a1 != a2 {
+					return false;
+				}
+			}
+			for (
+				Bond {
+					owner: o1,
+					amount: a1,
+				},
+				Bond {
+					owner: o2,
+					amount: a2,
+				},
+			) in self
+				.bottom_delegations
+				.iter()
+				.zip(other.bottom_delegations.iter())
+			{
+				if o1 != o2 || a1 != a2 {
+					return false;
+				}
+			}
+			true
+		}
+	}
+
 	/// Convey relevant information describing if a delegator was added to the top or bottom
 	/// Delegations added to the top yield a new total
 	#[derive(Clone, Copy, PartialEq, Encode, Decode, RuntimeDebug, TypeInfo)]
@@ -634,6 +691,35 @@ pub mod pallet {
 		pub status: DelegatorStatus,
 	}
 
+	// Temporary manual implementation for migration testing purposes
+	impl<A: PartialEq, B: PartialEq> PartialEq for Delegator<A, B> {
+		fn eq(&self, other: &Self) -> bool {
+			let must_be_true = self.id == other.id
+				&& self.total == other.total
+				&& self.requests == other.requests
+				&& self.status == other.status;
+			if !must_be_true {
+				return false;
+			}
+			for (
+				Bond {
+					owner: o1,
+					amount: a1,
+				},
+				Bond {
+					owner: o2,
+					amount: a2,
+				},
+			) in self.delegations.0.iter().zip(other.delegations.0.iter())
+			{
+				if o1 != o2 || a1 != a2 {
+					return false;
+				}
+			}
+			true
+		}
+	}
+
 	impl<
 			AccountId: Ord + Clone + Default,
 			Balance: Copy
@@ -1005,7 +1091,7 @@ pub mod pallet {
 		pub action: DelegationChange,
 	}
 
-	#[derive(Clone, Encode, Decode, RuntimeDebug, TypeInfo)]
+	#[derive(Clone, Encode, PartialEq, Decode, RuntimeDebug, TypeInfo)]
 	/// Pending requests to mutate delegations for each delegator
 	pub struct PendingDelegationRequests<AccountId, Balance> {
 		/// Number of pending revocations (necessary for determining whether revoke is exit)

pallets/parachain-staking/src/migrations.rs

@@ -20,6 +20,10 @@ use crate::{
 	BalanceOf, CandidateState, CollatorCandidate, CollatorState2, Config, DelegatorState,
 	ExitQueue2, NominatorState2, Points, Round, Staked,
 };
+#[cfg(feature = "try-runtime")]
+use crate::{Collator2, Delegator, Nominator2};
+#[cfg(feature = "try-runtime")]
+use frame_support::Twox64Concat;
 use frame_support::{
 	pallet_prelude::PhantomData,
 	traits::{Get, OnRuntimeUpgrade},
@@ -83,7 +87,10 @@ impl<T: Config> OnRuntimeUpgrade for RemoveExitQueue<T> {
 
 	#[cfg(feature = "try-runtime")]
 	fn pre_upgrade() -> Result<(), &'static str> {
-		use frame_support::{storage::migration::storage_iter, traits::OnRuntimeUpgradeHelpersExt};
+		use frame_support::{
+			storage::migration::{storage_iter, storage_key_iter},
+			traits::OnRuntimeUpgradeHelpersExt,
+		};
 
 		let pallet_prefix: &[u8] = b"ParachainStaking";
 		let collator_state_prefix: &[u8] = b"CollatorState2";
@@ -150,11 +157,12 @@ impl<T: Config> OnRuntimeUpgrade for RemoveExitQueue<T> {
 		// Check that our example candidate is converted correctly
 		if new_candidate_count > 0 {
 			let (account, original_collator_state): (
-				T::AuthorId,
+				T::AccountId,
 				Collator2<T::AccountId, BalanceOf<T>>,
 			) = Self::get_temp_storage("example_collator").expect("qed");
 			let new_candidate_state = CandidateState::<T>::get(account).expect("qed");
-			let old_candidate_converted: CollatorCandidate<_, _> = original_collator_state.into();
+			let old_candidate_converted: CollatorCandidate<T::AccountId, BalanceOf<T>> =
+				original_collator_state.into();
 			assert_eq!(new_candidate_state, old_candidate_converted);
 		}
 
@@ -164,15 +172,16 @@ impl<T: Config> OnRuntimeUpgrade for RemoveExitQueue<T> {
 		let new_delegator_count = DelegatorState::<T>::iter().count() as u64;
 		assert_eq!(old_nominator_count, new_delegator_count);
 
-		// Check that our example nominator is converted correctly
+		// Check that our example delegator is converted correctly
 		if new_delegator_count > 0 {
-			let (account, original_nominator_state): (
-				T::AuthorId,
+			let (account, original_delegator_state): (
+				T::AccountId,
 				Nominator2<T::AccountId, BalanceOf<T>>,
 			) = Self::get_temp_storage("example_nominator").expect("qed");
-			let new_candidate_state = DelegatorState::<T>::get(account).expect("qed");
-			let old_candidate_converted: Delegator<_, _> = original_nominator_state.into();
-			assert_eq!(old_candidate_converted, new_candidate_state);
+			let new_delegator_state = DelegatorState::<T>::get(&account).expect("qed");
+			let old_delegator_converted: Delegator<T::AccountId, BalanceOf<T>> =
+				migrate_nominator_to_delegator_state::<T>(account, original_delegator_state);
+			assert_eq!(old_delegator_converted, new_delegator_state);
 		}
 		Ok(())
 	}