From cadd20a4ff45724293ac3c01902348daae3716cc Mon Sep 17 00:00:00 2001
From: Xavier Lau <xavier@inv.cafe>
Date: Mon, 20 Nov 2023 00:55:52 +0800
Subject: [PATCH] Fix auto payout

---
 pallet/staking/src/lib.rs   |  4 +--
 pallet/staking/src/mock.rs  | 14 +++++------
 pallet/staking/src/tests.rs | 50 +++++++++++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 9 deletions(-)

diff --git a/pallet/staking/src/lib.rs b/pallet/staking/src/lib.rs
index 906c7293e..63007369b 100644
--- a/pallet/staking/src/lib.rs
+++ b/pallet/staking/src/lib.rs
@@ -357,7 +357,6 @@ pub mod pallet {
 			if !T::ShouldEndSession::get() {
 				call_on_exposure!(<Previous<T>>::iter_keys()
 					// TODO?: make this value adjustable
-					.drain()
 					.take(1)
 					.fold(Zero::zero(), |acc, e| acc
 						+ Self::payout_inner(e).unwrap_or(Zero::zero())))
@@ -902,6 +901,7 @@ pub mod pallet {
 
 				<PendingRewards<T>>::mutate(&c, |u| *u = u.map(|u| u + r).or(Some(r)));
 
+				// TODO: merge into one call
 				if T::InflationManager::reward(&staking_pot, r).is_ok() {
 					actual_reward += r;
 				} else {
@@ -924,7 +924,7 @@ pub mod pallet {
 		/// Pay the reward to the collator and its nominators.
 		pub fn payout_inner(collator: T::AccountId) -> Result<Weight, DispatchError> {
 			let c_exposure =
-				call_on_exposure!(<Previous<T>>::get(&collator).ok_or(<Error<T>>::NoReward)?)?;
+				call_on_exposure!(<Previous<T>>::take(&collator).ok_or(<Error<T>>::NoReward)?)?;
 			let c_total_payout =
 				<PendingRewards<T>>::take(&collator).ok_or(<Error<T>>::NoReward)?;
 			let mut c_payout = c_exposure.commission * c_total_payout;
diff --git a/pallet/staking/src/mock.rs b/pallet/staking/src/mock.rs
index b3be52197..9806e9499 100644
--- a/pallet/staking/src/mock.rs
+++ b/pallet/staking/src/mock.rs
@@ -401,9 +401,12 @@ impl Efflux {
 	}
 
 	pub fn block(number: BlockNumber) {
-		for _ in 0..number {
-			initialize_block(System::block_number() + 1)
-		}
+		let now = System::block_number();
+
+		(now..now + number).for_each(|n| {
+			initialize_block(n + 1);
+			finalize_block(n + 1);
+		});
 	}
 }
 
@@ -498,10 +501,7 @@ pub fn new_session() {
 	let now = System::block_number();
 	let target = now + <Period as sp_runtime::traits::Get<BlockNumber>>::get();
 
-	(now + 1..=target).for_each(|i| {
-		initialize_block(i);
-		finalize_block(i);
-	});
+	(now..target).for_each(|_| Efflux::block(1));
 }
 
 pub fn payout() {
diff --git a/pallet/staking/src/tests.rs b/pallet/staking/src/tests.rs
index 7ceb9124f..b9395cece 100644
--- a/pallet/staking/src/tests.rs
+++ b/pallet/staking/src/tests.rs
@@ -756,6 +756,56 @@ fn payout_should_work() {
 	});
 }
 
+#[test]
+fn auto_payout_should_work() {
+	ExtBuilder::default().collator_count(2).build().execute_with(|| {
+		(1..=2).for_each(|i| {
+			assert_ok!(Staking::stake(
+				RuntimeOrigin::signed(i),
+				0,
+				i as Balance * UNIT,
+				Vec::new()
+			));
+			assert_ok!(Staking::collect(RuntimeOrigin::signed(i), Perbill::from_percent(i * 10)));
+			assert_ok!(Staking::nominate(RuntimeOrigin::signed(i), i));
+		});
+		(3..=4).for_each(|i| {
+			assert_ok!(Staking::stake(
+				RuntimeOrigin::signed(i),
+				0,
+				(11 - i as Balance) * UNIT,
+				Vec::new()
+			));
+			assert_ok!(Staking::nominate(RuntimeOrigin::signed(i), i - 2));
+		});
+		new_session();
+		new_session();
+
+		Efflux::time(<Period as Get<u64>>::get() as Moment);
+		Staking::note_authors(&[1, 2]);
+		new_session();
+		(1..=4).for_each(|i| assert_eq!(Balances::free_balance(i), 1_000 * UNIT));
+
+		Efflux::block(1);
+		assert_eq!(Balances::free_balance(1), 1000000758879022790250);
+		assert_eq!(Balances::free_balance(2), 1000000000000000000000);
+		assert_eq!(Balances::free_balance(3), 1000003035516089643241);
+		assert_eq!(Balances::free_balance(4), 1000000000000000000000);
+
+		Efflux::block(1);
+		assert_eq!(Balances::free_balance(1), 1000000758879022790250);
+		assert_eq!(Balances::free_balance(2), 1000001433438163308069);
+		assert_eq!(Balances::free_balance(3), 1000003035516089643241);
+		assert_eq!(Balances::free_balance(4), 1000002360956952540378);
+
+		Efflux::block(1);
+		assert_eq!(Balances::free_balance(1), 1000000758879022790250);
+		assert_eq!(Balances::free_balance(2), 1000001433438163308069);
+		assert_eq!(Balances::free_balance(3), 1000003035516089643241);
+		assert_eq!(Balances::free_balance(4), 1000002360956952540378);
+	});
+}
+
 #[test]
 fn on_new_session_should_work() {
 	ExtBuilder::default().collator_count(2).genesis_collator().build().execute_with(|| {