feat: add native restaking test

pallets/multi-asset-delegation/src/functions/delegate.rs

@@ -441,7 +441,7 @@ impl<T: Config> Pallet<T> {
 		let nominated_amount = Self::verify_nomination_amount(&who, amount)?;
 
 		Delegators::<T>::try_mutate(&who, |maybe_metadata| -> DispatchResult {
-			let metadata = maybe_metadata.as_mut().ok_or(Error::<T>::NotDelegator)?;
+			let metadata = maybe_metadata.get_or_insert_with(Default::default);
 
 			// Calculate new total after this delegation
 			let current_total = metadata.total_nomination_delegations();

pallets/multi-asset-delegation/src/lib.rs

@@ -89,7 +89,7 @@ pub mod pallet {
 	use scale_info::TypeInfo;
 	use sp_core::H160;
 	use sp_runtime::traits::{MaybeSerializeDeserialize, Member, Zero};
-	use sp_staking::{SessionIndex, StakingInterface};
+	use sp_staking::{OnStakingUpdate, SessionIndex, StakingInterface};
 	use sp_std::{fmt::Debug, prelude::*, vec::Vec};
 	use tangle_primitives::traits::RewardsManager;
 	use tangle_primitives::types::rewards::LockMultiplier;
@@ -1230,4 +1230,6 @@ pub mod pallet {
 			I::end_session(i)
 		}
 	}
+
+	impl<C: Config> OnStakingUpdate<C::AccountId, BalanceOf<C>> for Pallet<C> {}
 }

pallets/multi-asset-delegation/src/mock.rs

@@ -220,7 +220,7 @@ impl pallet_staking::Config for Runtime {
 	type TargetList = pallet_staking::UseValidatorsMap<Self>;
 	type MaxUnlockingChunks = ConstU32<32>;
 	type HistoryDepth = ConstU32<84>;
-	type EventListeners = ();
+	type EventListeners = MultiAssetDelegation;
 	type BenchmarkingConfig = pallet_staking::TestBenchmarkingConfig;
 	type NominationsQuota = pallet_staking::FixedNominationsQuota<MAX_QUOTA_NOMINATIONS>;
 	type WeightInfo = ();
@@ -461,6 +461,7 @@ pub fn new_test_ext() -> sp_io::TestExternalities {
 	ext
 }
 
+pub const TNT: AssetId = 0;
 pub const USDC_ERC20: H160 = H160([0x23; 20]);
 pub const VDOT: AssetId = 4;
 
@@ -519,6 +520,14 @@ pub fn new_test_ext_raw_authorities() -> sp_io::TestExternalities {
 
 	evm_config.assimilate_storage(&mut t).unwrap();
 
+	let staking_config = pallet_staking::GenesisConfig::<Runtime> {
+		validator_count: 3,
+		invulnerables: authorities.clone(),
+		..Default::default()
+	};
+
+	staking_config.assimilate_storage(&mut t).unwrap();
+
 	// assets_config.assimilate_storage(&mut t).unwrap();
 	let mut ext = sp_io::TestExternalities::new(t);
 	ext.register_extension(KeystoreExt(Arc::new(MemoryKeystore::new()) as KeystorePtr));

pallets/multi-asset-delegation/src/tests/delegate.rs

@@ -16,8 +16,11 @@
 #![allow(clippy::all)]
 use super::*;
 use crate::{CurrentRound, Error};
-use frame_support::{assert_noop, assert_ok};
-use sp_keyring::AccountKeyring::{Alice, Bob};
+use frame_support::{
+	assert_noop, assert_ok,
+	traits::{OnFinalize, OnInitialize},
+};
+use sp_keyring::AccountKeyring::{Alice, Bob, Dave};
 use tangle_primitives::services::Asset;
 
 #[test]
@@ -119,6 +122,17 @@ fn schedule_delegator_unstake_should_work() {
 		assert_eq!(request.asset_id, asset_id);
 		assert_eq!(request.amount, amount);
 
+		// Check the operator metadata
+		let operator_metadata = MultiAssetDelegation::operator_info(operator.clone()).unwrap();
+		assert_eq!(operator_metadata.delegation_count, 1);
+		assert_eq!(operator_metadata.delegations.len(), 1);
+		// Move to next round
+		CurrentRound::<Runtime>::put(10);
+		// Execute the unstake
+		assert_ok!(MultiAssetDelegation::execute_delegator_unstake(RuntimeOrigin::signed(
+			who.clone()
+		),));
+
 		// Check the operator metadata
 		let operator_metadata = MultiAssetDelegation::operator_info(operator.clone()).unwrap();
 		assert_eq!(operator_metadata.delegation_count, 0);
@@ -231,8 +245,8 @@ fn cancel_delegator_unstake_should_work() {
 
 		// Check the operator metadata
 		let operator_metadata = MultiAssetDelegation::operator_info(operator.clone()).unwrap();
-		assert_eq!(operator_metadata.delegation_count, 0);
-		assert_eq!(operator_metadata.delegations.len(), 0);
+		assert_eq!(operator_metadata.delegation_count, 1);
+		assert_eq!(operator_metadata.delegations.len(), 1);
 
 		assert_ok!(MultiAssetDelegation::cancel_delegator_unstake(
 			RuntimeOrigin::signed(who.clone()),
@@ -310,7 +324,7 @@ fn cancel_delegator_unstake_should_update_already_existing() {
 		assert_eq!(operator_metadata.delegations.len(), 1);
 		let operator_delegation = &operator_metadata.delegations[0];
 		assert_eq!(operator_delegation.delegator, who.clone());
-		assert_eq!(operator_delegation.amount, amount - 10);
+		assert_eq!(operator_delegation.amount, amount);
 		assert_eq!(operator_delegation.asset_id, asset_id);
 
 		assert_ok!(MultiAssetDelegation::cancel_delegator_unstake(
@@ -539,3 +553,63 @@ fn delegate_should_not_create_multiple_on_repeat_delegation() {
 		assert_eq!(updated_operator_delegation.asset_id, asset_id);
 	});
 }
+
+#[test]
+fn native_restaking_should_work() {
+	new_test_ext().execute_with(|| {
+		// Arrange
+		let who: AccountId = Dave.into();
+		let validator = Staking::invulnerables()[0].clone();
+		let operator: AccountId = Alice.into();
+		let amount = 100_000;
+		let delegate_amount = amount / 2;
+		// Bond Some TNT
+		assert_ok!(Staking::bond(
+			RuntimeOrigin::signed(who.clone()),
+			amount,
+			pallet_staking::RewardDestination::Staked
+		));
+		// Nominate the validator
+		assert_ok!(Staking::nominate(RuntimeOrigin::signed(who.clone()), vec![validator.clone()]));
+
+		System::set_block_number(2);
+		Session::on_initialize(2);
+		Staking::on_initialize(2);
+		Session::on_finalize(2);
+		Staking::on_finalize(2);
+		// Assert
+		let ledger = Staking::ledger(sp_staking::StakingAccount::Stash(who.clone())).unwrap();
+		assert_eq!(ledger.active, amount);
+		assert_eq!(ledger.total, amount);
+		assert_eq!(ledger.unlocking.len(), 0);
+
+		assert_ok!(MultiAssetDelegation::join_operators(
+			RuntimeOrigin::signed(operator.clone()),
+			10_000
+		));
+
+		// Restake
+		assert_ok!(MultiAssetDelegation::delegate_nomination(
+			RuntimeOrigin::signed(who.clone()),
+			operator.clone(),
+			delegate_amount,
+			Default::default(),
+		));
+		// Assert
+		let metadata = MultiAssetDelegation::delegators(who.clone()).unwrap();
+		assert_eq!(metadata.delegations.len(), 1);
+		let delegation = &metadata.delegations[0];
+		assert_eq!(delegation.operator, operator.clone());
+		assert_eq!(delegation.amount, delegate_amount);
+		assert_eq!(delegation.asset_id, Asset::Custom(TNT));
+		// Check the locks
+		let locks = pallet_balances::Pallet::<Runtime>::locks(&who);
+		// 1 lock for the staking
+		// 1 lock for the delegation
+		assert_eq!(locks.len(), 2);
+		assert_eq!(&locks[0].id, b"staking ");
+		assert_eq!(locks[0].amount, amount);
+		assert_eq!(&locks[1].id, b"delegate");
+		assert_eq!(locks[1].amount, delegate_amount);
+	});
+}

pallets/multi-asset-delegation/src/tests/session_manager.rs

@@ -160,7 +160,7 @@ fn handle_round_change_with_unstake_should_work() {
 		assert_eq!(snapshot1.stake, 10_000);
 		assert_eq!(snapshot1.delegations.len(), 1);
 		assert_eq!(snapshot1.delegations[0].delegator, delegator1.clone());
-		assert_eq!(snapshot1.delegations[0].amount, amount1 - unstake_amount); // Amount reduced by unstake_amount
+		assert_eq!(snapshot1.delegations[0].amount, amount1); // Amount should be the same
 		assert_eq!(snapshot1.delegations[0].asset_id, asset_id);
 
 		// Check the snapshot for operator2