Fix auto payout

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;

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() {

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(|| {