Allow Staking tests to run with session length other than 1 (#7719)

* fix periodic session

* Allow staking tests to run with session lengths other than 1.

* Update frame/staking/src/mock.rs

Co-authored-by: Guillaume Thiolliere <gui.thiolliere@gmail.com>

* Fix all tests with session length 5.

* Test for active != current

* Better doc

* Update frame/staking/src/lib.rs

Co-authored-by: Guillaume Thiolliere <gui.thiolliere@gmail.com>

* also set the timestamp properly.

* trigger CI

* Revert "trigger CI"

This reverts commit 0f254944cdad848aa6e63bd8a618db95447a8e68.

* Update frame/staking/src/lib.rs

Co-authored-by: Guillaume Thiolliere <gui.thiolliere@gmail.com>

substrate/frame/session/src/lib.rs

@@ -17,8 +17,8 @@
 
 //! # Session Module
 //!
-//! The Session module allows validators to manage their session keys, provides a function for changing
-//! the session length, and handles session rotation.
+//! The Session module allows validators to manage their session keys, provides a function for
+//! changing the session length, and handles session rotation.
 //!
 //! - [`session::Config`](./trait.Config.html)
 //! - [`Call`](./enum.Call.html)
@@ -29,34 +29,39 @@
 //! ### Terminology
 //! <!-- Original author of paragraph: @gavofyork -->
 //!
-//! - **Session:** A session is a period of time that has a constant set of validators. Validators can only join
-//! or exit the validator set at a session change. It is measured in block numbers. The block where a session is
-//! ended is determined by the `ShouldEndSession` trait. When the session is ending, a new validator set
-//! can be chosen by `OnSessionEnding` implementations.
-//! - **Session key:** A session key is actually several keys kept together that provide the various signing
-//! functions required by network authorities/validators in pursuit of their duties.
-//! - **Validator ID:** Every account has an associated validator ID. For some simple staking systems, this
-//! may just be the same as the account ID. For staking systems using a stash/controller model,
-//! the validator ID would be the stash account ID of the controller.
+//! - **Session:** A session is a period of time that has a constant set of validators. Validators
+//!   can only join or exit the validator set at a session change. It is measured in block numbers.
+//!   The block where a session is ended is determined by the `ShouldEndSession` trait. When the
+//!   session is ending, a new validator set can be chosen by `OnSessionEnding` implementations.
+//!
+//! - **Session key:** A session key is actually several keys kept together that provide the various
+//!   signing functions required by network authorities/validators in pursuit of their duties.
+//! - **Validator ID:** Every account has an associated validator ID. For some simple staking
+//!   systems, this may just be the same as the account ID. For staking systems using a
+//!   stash/controller model, the validator ID would be the stash account ID of the controller.
+//!
 //! - **Session key configuration process:** Session keys are set using `set_keys` for use not in
-//! the next session, but the session after next. They are stored in `NextKeys`, a mapping between
-//! the caller's `ValidatorId` and the session keys provided. `set_keys` allows users to set their
-//! session key prior to being selected as validator.
-//! It is a public call since it uses `ensure_signed`, which checks that the origin is a signed account.
-//! As such, the account ID of the origin stored in `NextKeys` may not necessarily be associated with
-//! a block author or a validator. The session keys of accounts are removed once their account balance is zero.
+//!   the next session, but the session after next. They are stored in `NextKeys`, a mapping between
+//!   the caller's `ValidatorId` and the session keys provided. `set_keys` allows users to set their
+//!   session key prior to being selected as validator. It is a public call since it uses
+//!   `ensure_signed`, which checks that the origin is a signed account. As such, the account ID of
+//!   the origin stored in `NextKeys` may not necessarily be associated with a block author or a
+//!   validator. The session keys of accounts are removed once their account balance is zero.
+//!
 //! - **Session length:** This pallet does not assume anything about the length of each session.
-//! Rather, it relies on an implementation of `ShouldEndSession` to dictate a new session's start.
-//! This pallet provides the `PeriodicSessions` struct for simple periodic sessions.
-//! - **Session rotation configuration:** Configure as either a 'normal' (rewardable session where rewards are
-//! applied) or 'exceptional' (slashable) session rotation.
+//!   Rather, it relies on an implementation of `ShouldEndSession` to dictate a new session's start.
+//!   This pallet provides the `PeriodicSessions` struct for simple periodic sessions.
+//!
+//! - **Session rotation configuration:** Configure as either a 'normal' (rewardable session where
+//!   rewards are applied) or 'exceptional' (slashable) session rotation.
+//!
 //! - **Session rotation process:** At the beginning of each block, the `on_initialize` function
-//! queries the provided implementation of `ShouldEndSession`. If the session is to end the newly
-//! activated validator IDs and session keys are taken from storage and passed to the
-//! `SessionHandler`. The validator set supplied by `SessionManager::new_session` and the corresponding session
-//! keys, which may have been registered via `set_keys` during the previous session, are written
-//! to storage where they will wait one session before being passed to the `SessionHandler`
-//! themselves.
+//!   queries the provided implementation of `ShouldEndSession`. If the session is to end the newly
+//!   activated validator IDs and session keys are taken from storage and passed to the
+//!   `SessionHandler`. The validator set supplied by `SessionManager::new_session` and the
+//!   corresponding session keys, which may have been registered via `set_keys` during the previous
+//!   session, are written to storage where they will wait one session before being passed to the
+//!   `SessionHandler` themselves.
 //!
 //! ### Goals
 //!
@@ -75,7 +80,7 @@
 //! ### Public Functions
 //!
 //! - `rotate_session` - Change to the next session. Register the new authority set. Queue changes
-//! for next session rotation.
+//!   for next session rotation.
 //! - `disable_index` - Disable a validator by index.
 //! - `disable` - Disable a validator by Validator ID
 //!
@@ -83,13 +88,14 @@
 //!
 //! ### Example from the FRAME
 //!
-//! The [Staking pallet](../pallet_staking/index.html) uses the Session pallet to get the validator set.
+//! The [Staking pallet](../pallet_staking/index.html) uses the Session pallet to get the validator
+//! set.
 //!
 //! ```
 //! use pallet_session as session;
 //!
 //! fn validators<T: pallet_session::Config>() -> Vec<<T as pallet_session::Config>::ValidatorId> {
-//!	<pallet_session::Module<T>>::validators()
+//! <pallet_session::Module<T>>::validators()
 //! }
 //! # fn main(){}
 //! ```
@@ -166,7 +172,7 @@ impl<
 					period.saturating_sub(block_after_last_session)
 				)
 			} else {
-				Zero::zero()
+				now
 			}
 		} else {
 			offset
@@ -174,10 +180,10 @@ impl<
 	}
 
 	fn weight(_now: BlockNumber) -> Weight {
-		// Weight note: `estimate_next_session_rotation` has no storage reads and trivial computational overhead.
-		// There should be no risk to the chain having this weight value be zero for now.
-		// However, this value of zero was not properly calculated, and so it would be reasonable
-		// to come back here and properly calculate the weight of this function.
+		// Weight note: `estimate_next_session_rotation` has no storage reads and trivial
+		// computational overhead. There should be no risk to the chain having this weight value be
+		// zero for now. However, this value of zero was not properly calculated, and so it would be
+		// reasonable to come back here and properly calculate the weight of this function.
 		0
 	}
 }
@@ -186,17 +192,17 @@ impl<
 pub trait SessionManager<ValidatorId> {
 	/// Plan a new session, and optionally provide the new validator set.
 	///
-	/// Even if the validator-set is the same as before, if any underlying economic
-	/// conditions have changed (i.e. stake-weights), the new validator set must be returned.
-	/// This is necessary for consensus engines making use of the session module to
-	/// issue a validator-set change so misbehavior can be provably associated with the new
-	/// economic conditions as opposed to the old.
-	/// The returned validator set, if any, will not be applied until `new_index`.
-	/// `new_index` is strictly greater than from previous call.
+	/// Even if the validator-set is the same as before, if any underlying economic conditions have
+	/// changed (i.e. stake-weights), the new validator set must be returned. This is necessary for
+	/// consensus engines making use of the session module to issue a validator-set change so
+	/// misbehavior can be provably associated with the new economic conditions as opposed to the
+	/// old. The returned validator set, if any, will not be applied until `new_index`. `new_index`
+	/// is strictly greater than from previous call.
 	///
 	/// The first session start at index 0.
 	///
-	/// `new_session(session)` is guaranteed to be called before `end_session(session-1)`.
+	/// `new_session(session)` is guaranteed to be called before `end_session(session-1)`. In other
+	/// words, a new session must always be planned before an ongoing one can be finished.
 	fn new_session(new_index: SessionIndex) -> Option<Vec<ValidatorId>>;
 	/// End the session.
 	///
@@ -205,7 +211,7 @@ pub trait SessionManager<ValidatorId> {
 	fn end_session(end_index: SessionIndex);
 	/// Start the session.
 	///
-	/// The session start to be used for validation
+	/// The session start to be used for validation.
 	fn start_session(start_index: SessionIndex);
 }
 
@@ -242,7 +248,7 @@ pub trait SessionHandler<ValidatorId> {
 
 	/// A notification for end of the session.
 	///
-	/// Note it is triggered before any `SessionManager::end_session` handlers,
+	/// Note it is triggered before any [`SessionManager::end_session`] handlers,
 	/// so we can still affect the validator set.
 	fn on_before_session_ending() {}
 

substrate/frame/session/src/tests.rs

@@ -252,31 +252,32 @@ fn session_changed_flag_works() {
 
 #[test]
 fn periodic_session_works() {
-	struct Period;
-	struct Offset;
 
-	impl Get<u64> for Period {
-		fn get() -> u64 { 10 }
+	frame_support::parameter_types! {
+		const Period: u64 = 10;
+		const Offset: u64 = 3;
 	}
 
-	impl Get<u64> for Offset {
-		fn get() -> u64 { 3 }
-	}
-
-
 	type P = PeriodicSessions<Period, Offset>;
 
-	for i in 0..3 {
+	for i in 0u64..3 {
 		assert!(!P::should_end_session(i));
+		assert_eq!(P::estimate_next_session_rotation(i).unwrap(), 3);
 	}
 
-	assert!(P::should_end_session(3));
+	assert!(P::should_end_session(3u64));
+	assert_eq!(P::estimate_next_session_rotation(3u64).unwrap(), 3);
 
-	for i in (1..10).map(|i| 3 + i) {
+	for i in (1u64..10).map(|i| 3 + i) {
 		assert!(!P::should_end_session(i));
+		assert_eq!(P::estimate_next_session_rotation(i).unwrap(), 13);
 	}
 
-	assert!(P::should_end_session(13));
+	assert!(P::should_end_session(13u64));
+	assert_eq!(P::estimate_next_session_rotation(13u64).unwrap(), 13);
+
+	assert!(!P::should_end_session(14u64));
+	assert_eq!(P::estimate_next_session_rotation(14u64).unwrap(), 23);
 }
 
 #[test]

substrate/frame/staking/src/lib.rs

@@ -952,11 +952,14 @@ decl_storage! {
 
 		/// The active era information, it holds index and start.
 		///
-		/// The active era is the era currently rewarded.
-		/// Validator set of this era must be equal to `SessionInterface::validators`.
+		/// The active era is the era being currently rewarded. Validator set of this era must be
+		/// equal to [`SessionInterface::validators`].
 		pub ActiveEra get(fn active_era): Option<ActiveEraInfo>;
 
 		/// The session index at which the era start for the last `HISTORY_DEPTH` eras.
+		///
+		/// Note: This tracks the starting session (i.e. session index when era start being active)
+		/// for the eras in `[CurrentEra - HISTORY_DEPTH, CurrentEra]`.
 		pub ErasStartSessionIndex get(fn eras_start_session_index):
 			map hasher(twox_64_concat) EraIndex => Option<SessionIndex>;
 
@@ -2630,14 +2633,17 @@ impl<T: Config> Module<T> {
 	/// Start a session potentially starting an era.
 	fn start_session(start_session: SessionIndex) {
 		let next_active_era = Self::active_era().map(|e| e.index + 1).unwrap_or(0);
+		// This is only `Some` when current era has already progressed to the next era, while the
+		// active era is one behind (i.e. in the *last session of the active era*, or *first session
+		// of the new current era*, depending on how you look at it).
 		if let Some(next_active_era_start_session_index) =
 			Self::eras_start_session_index(next_active_era)
 		{
 			if next_active_era_start_session_index == start_session {
 				Self::start_era(start_session);
 			} else if next_active_era_start_session_index < start_session {
-				// This arm should never happen, but better handle it than to stall the
-				// staking pallet.
+				// This arm should never happen, but better handle it than to stall the staking
+				// pallet.
 				frame_support::print("Warning: A session appears to have been skipped.");
 				Self::start_era(start_session);
 			}
@@ -2893,9 +2899,11 @@ impl<T: Config> Module<T> {
 	/// Self votes are added and nominations before the most recent slashing span are ignored.
 	///
 	/// No storage item is updated.
-	pub fn do_phragmen<Accuracy: PerThing>(iterations: usize)
-	-> Option<PrimitiveElectionResult<T::AccountId, Accuracy>>
-		where ExtendedBalance: From<InnerOf<Accuracy>>
+	pub fn do_phragmen<Accuracy: PerThing>(
+		iterations: usize,
+	) -> Option<PrimitiveElectionResult<T::AccountId, Accuracy>>
+	where
+		ExtendedBalance: From<InnerOf<Accuracy>>,
 	{
 		let weight_of = Self::slashable_balance_of_fn();
 		let mut all_nominators: Vec<(T::AccountId, VoteWeight, Vec<T::AccountId>)> = Vec::new();
@@ -2928,7 +2936,11 @@ impl<T: Config> Module<T> {
 
 		if all_validators.len() < Self::minimum_validator_count().max(1) as usize {
 			// If we don't have enough candidates, nothing to do.
-			log!(error, "💸 Chain does not have enough staking candidates to operate. Era {:?}.", Self::current_era());
+			log!(
+				warn,
+				"💸 Chain does not have enough staking candidates to operate. Era {:?}.",
+				Self::current_era()
+			);
 			None
 		} else {
 			seq_phragmen::<_, Accuracy>(
@@ -3090,12 +3102,30 @@ impl<T: Config> Module<T> {
 /// some session can lag in between the newest session planned and the latest session started.
 impl<T: Config> pallet_session::SessionManager<T::AccountId> for Module<T> {
 	fn new_session(new_index: SessionIndex) -> Option<Vec<T::AccountId>> {
+		frame_support::debug::native::trace!(
+			target: LOG_TARGET,
+			"[{}] planning new_session({})",
+			<frame_system::Module<T>>::block_number(),
+			new_index
+		);
 		Self::new_session(new_index)
 	}
 	fn start_session(start_index: SessionIndex) {
+		frame_support::debug::native::trace!(
+			target: LOG_TARGET,
+			"[{}] starting start_session({})",
+			<frame_system::Module<T>>::block_number(),
+			start_index
+		);
 		Self::start_session(start_index)
 	}
 	fn end_session(end_index: SessionIndex) {
+		frame_support::debug::native::trace!(
+			target: LOG_TARGET,
+			"[{}] ending end_session({})",
+			<frame_system::Module<T>>::block_number(),
+			end_index
+		);
 		Self::end_session(end_index)
 	}
 }

substrate/frame/staking/src/mock.rs

@@ -38,6 +38,7 @@ use sp_staking::offence::{OffenceDetails, OnOffenceHandler};
 use std::{cell::RefCell, collections::HashSet};
 
 pub const INIT_TIMESTAMP: u64 = 30_000;
+pub const BLOCK_TIME: u64 = 1000;
 
 /// The AccountId alias in this test module.
 pub(crate) type AccountId = u64;
@@ -135,10 +136,11 @@ parameter_types! {
 		);
 	pub const MaxLocks: u32 = 1024;
 	pub static SessionsPerEra: SessionIndex = 3;
-	pub static ExistentialDeposit: Balance = 0;
+	pub static ExistentialDeposit: Balance = 1;
 	pub static SlashDeferDuration: EraIndex = 0;
 	pub static ElectionLookahead: BlockNumber = 0;
-	pub static Period: BlockNumber = 1;
+	pub static Period: BlockNumber = 5;
+	pub static Offset: BlockNumber = 0;
 	pub static MaxIterations: u32 = 0;
 }
 
@@ -175,7 +177,6 @@ impl pallet_balances::Config for Test {
 	type WeightInfo = ();
 }
 parameter_types! {
-	pub const Offset: BlockNumber = 0;
 	pub const UncleGenerations: u64 = 0;
 	pub const DisabledValidatorsThreshold: Perbill = Perbill::from_percent(25);
 }
@@ -286,46 +287,36 @@ where
 pub type Extrinsic = TestXt<Call, ()>;
 
 pub struct ExtBuilder {
-	session_length: BlockNumber,
-	election_lookahead: BlockNumber,
-	session_per_era: SessionIndex,
-	existential_deposit: Balance,
 	validator_pool: bool,
 	nominate: bool,
 	validator_count: u32,
 	minimum_validator_count: u32,
-	slash_defer_duration: EraIndex,
 	fair: bool,
 	num_validators: Option<u32>,
 	invulnerables: Vec<AccountId>,
 	has_stakers: bool,
-	max_offchain_iterations: u32,
+	initialize_first_session: bool,
 }
 
 impl Default for ExtBuilder {
 	fn default() -> Self {
 		Self {
-			session_length: 1,
-			election_lookahead: 0,
-			session_per_era: 3,
-			existential_deposit: 1,
 			validator_pool: false,
 			nominate: true,
 			validator_count: 2,
 			minimum_validator_count: 0,
-			slash_defer_duration: 0,
 			fair: true,
 			num_validators: None,
 			invulnerables: vec![],
 			has_stakers: true,
-			max_offchain_iterations: 0,
+			initialize_first_session: true,
 		}
 	}
 }
 
 impl ExtBuilder {
-	pub fn existential_deposit(mut self, existential_deposit: Balance) -> Self {
-		self.existential_deposit = existential_deposit;
+	pub fn existential_deposit(self, existential_deposit: Balance) -> Self {
+		EXISTENTIAL_DEPOSIT.with(|v| *v.borrow_mut() = existential_deposit);
 		self
 	}
 	pub fn validator_pool(mut self, validator_pool: bool) -> Self {
@@ -344,8 +335,8 @@ impl ExtBuilder {
 		self.minimum_validator_count = count;
 		self
 	}
-	pub fn slash_defer_duration(mut self, eras: EraIndex) -> Self {
-		self.slash_defer_duration = eras;
+	pub fn slash_defer_duration(self, eras: EraIndex) -> Self {
+		SLASH_DEFER_DURATION.with(|v| *v.borrow_mut() = eras);
 		self
 	}
 	pub fn fair(mut self, is_fair: bool) -> Self {
@@ -360,46 +351,43 @@ impl ExtBuilder {
 		self.invulnerables = invulnerables;
 		self
 	}
-	pub fn session_per_era(mut self, length: SessionIndex) -> Self {
-		self.session_per_era = length;
+	pub fn session_per_era(self, length: SessionIndex) -> Self {
+		SESSIONS_PER_ERA.with(|v| *v.borrow_mut() = length);
 		self
 	}
-	pub fn election_lookahead(mut self, look: BlockNumber) -> Self {
-		self.election_lookahead = look;
+	pub fn election_lookahead(self, look: BlockNumber) -> Self {
+		ELECTION_LOOKAHEAD.with(|v| *v.borrow_mut() = look);
 		self
 	}
-	pub fn session_length(mut self, length: BlockNumber) -> Self {
-		self.session_length = length;
+	pub fn period(self, length: BlockNumber) -> Self {
+		PERIOD.with(|v| *v.borrow_mut() = length);
 		self
 	}
 	pub fn has_stakers(mut self, has: bool) -> Self {
 		self.has_stakers = has;
 		self
 	}
-	pub fn max_offchain_iterations(mut self, iterations: u32) -> Self {
-		self.max_offchain_iterations = iterations;
+	pub fn max_offchain_iterations(self, iterations: u32) -> Self {
+		MAX_ITERATIONS.with(|v| *v.borrow_mut() = iterations);
 		self
 	}
 	pub fn offchain_election_ext(self) -> Self {
-		self.session_per_era(4)
-			.session_length(5)
-			.election_lookahead(3)
+		self.session_per_era(4).period(5).election_lookahead(3)
 	}
-	pub fn set_associated_constants(&self) {
-		EXISTENTIAL_DEPOSIT.with(|v| *v.borrow_mut() = self.existential_deposit);
-		SLASH_DEFER_DURATION.with(|v| *v.borrow_mut() = self.slash_defer_duration);
-		SESSIONS_PER_ERA.with(|v| *v.borrow_mut() = self.session_per_era);
-		ELECTION_LOOKAHEAD.with(|v| *v.borrow_mut() = self.election_lookahead);
-		PERIOD.with(|v| *v.borrow_mut() = self.session_length);
-		MAX_ITERATIONS.with(|v| *v.borrow_mut() = self.max_offchain_iterations);
+	pub fn initialize_first_session(mut self, init: bool) -> Self {
+		self.initialize_first_session = init;
+		self
+	}
+	pub fn offset(self, offset: BlockNumber) -> Self {
+		OFFSET.with(|v| *v.borrow_mut() = offset);
+		self
 	}
 	pub fn build(self) -> sp_io::TestExternalities {
 		sp_tracing::try_init_simple();
-		self.set_associated_constants();
 		let mut storage = frame_system::GenesisConfig::default()
 			.build_storage::<Test>()
 			.unwrap();
-		let balance_factor = if self.existential_deposit > 1 {
+		let balance_factor = if ExistentialDeposit::get() > 1 {
 			256
 		} else {
 			1
@@ -475,13 +463,17 @@ impl ExtBuilder {
 			SESSION.with(|x| *x.borrow_mut() = (validators.clone(), HashSet::new()));
 		});
 
-		// We consider all test to start after timestamp is initialized
-		// This must be ensured by having `timestamp::on_initialize` called before
-		// `staking::on_initialize`
-		ext.execute_with(|| {
-			System::set_block_number(1);
-			Timestamp::set_timestamp(INIT_TIMESTAMP);
-		});
+		if self.initialize_first_session {
+			// We consider all test to start after timestamp is initialized This must be ensured by
+			// having `timestamp::on_initialize` called before `staking::on_initialize`. Also, if
+			// session length is 1, then it is already triggered.
+			ext.execute_with(|| {
+				System::set_block_number(1);
+				Session::on_initialize(1);
+				Staking::on_initialize(1);
+				Timestamp::set_timestamp(INIT_TIMESTAMP);
+			});
+		}
 
 		ext
 	}
@@ -498,20 +490,12 @@ pub type Session = pallet_session::Module<Test>;
 pub type Timestamp = pallet_timestamp::Module<Test>;
 pub type Staking = Module<Test>;
 
-pub(crate) fn current_era() -> EraIndex {
-	Staking::current_era().unwrap()
-}
-
 fn post_conditions() {
 	check_nominators();
 	check_exposures();
 	check_ledgers();
 }
 
-pub(crate) fn active_era() -> EraIndex {
-	Staking::active_era().unwrap().index
-}
-
 fn check_ledgers() {
 	// check the ledger of all stakers.
 	Bonded::<Test>::iter().for_each(|(_, ctrl)| assert_ledger_consistent(ctrl))
@@ -593,6 +577,14 @@ fn assert_ledger_consistent(ctrl: AccountId) {
 	);
 }
 
+pub(crate) fn active_era() -> EraIndex {
+	Staking::active_era().unwrap().index
+}
+
+pub(crate) fn current_era() -> EraIndex {
+	Staking::current_era().unwrap()
+}
+
 pub(crate) fn bond_validator(stash: AccountId, ctrl: AccountId, val: Balance) {
 	let _ = Balances::make_free_balance_be(&stash, val);
 	let _ = Balances::make_free_balance_be(&ctrl, val);
@@ -625,52 +617,98 @@ pub(crate) fn bond_nominator(
 	assert_ok!(Staking::nominate(Origin::signed(ctrl), target));
 }
 
+/// Progress to the given block, triggering session and era changes as we progress.
+///
+/// This will finalize the previous block, initialize up to the given block, essentially simulating
+/// a block import/propose process where we first initialize the block, then execute some stuff (not
+/// in the function), and then finalize the block.
 pub(crate) fn run_to_block(n: BlockNumber) {
 	Staking::on_finalize(System::block_number());
-	for b in System::block_number() + 1..=n {
+	for b in (System::block_number() + 1)..=n {
 		System::set_block_number(b);
 		Session::on_initialize(b);
 		Staking::on_initialize(b);
+		Timestamp::set_timestamp(System::block_number() * BLOCK_TIME + INIT_TIMESTAMP);
 		if b != n {
 			Staking::on_finalize(System::block_number());
 		}
 	}
 }
 
+/// Progresses from the current block number (whatever that may be) to the `P * session_index + 1`.
+pub(crate) fn start_session(session_index: SessionIndex) {
+	let end: u64 = if Offset::get().is_zero() {
+		(session_index as u64) * Period::get()
+	} else {
+		Offset::get() + (session_index.saturating_sub(1) as u64) * Period::get()
+	};
+	run_to_block(end);
+	// session must have progressed properly.
+	assert_eq!(
+		Session::current_index(),
+		session_index,
+		"current session index = {}, expected = {}",
+		Session::current_index(),
+		session_index,
+	);
+}
+
+/// Go one session forward.
 pub(crate) fn advance_session() {
 	let current_index = Session::current_index();
 	start_session(current_index + 1);
 }
 
-pub(crate) fn start_session(session_index: SessionIndex) {
-	assert_eq!(<Period as Get<BlockNumber>>::get(), 1, "start_session can only be used with session length 1.");
-	for i in Session::current_index()..session_index {
-		Staking::on_finalize(System::block_number());
-		System::set_block_number((i + 1).into());
-		Timestamp::set_timestamp(System::block_number() * 1000 + INIT_TIMESTAMP);
-		Session::on_initialize(System::block_number());
-		Staking::on_initialize(System::block_number());
-	}
-
-	assert_eq!(Session::current_index(), session_index);
-}
-
-// This start and activate the era given.
-// Because the mock use pallet-session which delays session by one, this will be one session after
-// the election happened, not the first session after the election has happened.
-pub(crate) fn start_era(era_index: EraIndex) {
+/// Progress until the given era.
+pub(crate) fn start_active_era(era_index: EraIndex) {
 	start_session((era_index * <SessionsPerEra as Get<u32>>::get()).into());
-	assert_eq!(Staking::current_era().unwrap(), era_index);
-	assert_eq!(Staking::active_era().unwrap().index, era_index);
+	assert_eq!(active_era(), era_index);
+	// One way or another, current_era must have changed before the active era, so they must match
+	// at this point.
+	assert_eq!(current_era(), active_era());
 }
 
 pub(crate) fn current_total_payout_for_duration(duration: u64) -> Balance {
-	inflation::compute_total_payout(
+	let reward = inflation::compute_total_payout(
 		<Test as Config>::RewardCurve::get(),
-		Staking::eras_total_stake(Staking::active_era().unwrap().index),
+		Staking::eras_total_stake(active_era()),
 		Balances::total_issuance(),
 		duration,
-	).0
+	)
+	.0;
+	assert!(reward > 0);
+	reward
+}
+
+pub(crate) fn maximum_payout_for_duration(duration: u64) -> Balance {
+	inflation::compute_total_payout(
+		<Test as Config>::RewardCurve::get(),
+		0,
+		Balances::total_issuance(),
+		duration,
+	)
+	.1
+}
+
+/// Time it takes to finish a session.
+///
+/// Note, if you see `time_per_session() - BLOCK_TIME`, it is fine. This is because we set the
+/// timestamp after on_initialize, so the timestamp is always one block old.
+pub(crate) fn time_per_session() -> u64 {
+	Period::get() * BLOCK_TIME
+}
+
+/// Time it takes to finish an era.
+///
+/// Note, if you see `time_per_era() - BLOCK_TIME`, it is fine. This is because we set the
+/// timestamp after on_initialize, so the timestamp is always one block old.
+pub(crate) fn time_per_era() -> u64 {
+	time_per_session() * SessionsPerEra::get() as u64
+}
+
+/// Time that will be calculated for the reward per era.
+pub(crate) fn reward_time_per_era() -> u64 {
+	time_per_era() - BLOCK_TIME
 }
 
 pub(crate) fn reward_all_elected() {
@@ -939,8 +977,11 @@ pub(crate) fn make_all_reward_payment(era: EraIndex) {
 	// reward validators
 	for validator_controller in validators_with_reward.iter().filter_map(Staking::bonded) {
 		let ledger = <Ledger<Test>>::get(&validator_controller).unwrap();
-
-		assert_ok!(Staking::payout_stakers(Origin::signed(1337), ledger.stash, era));
+		assert_ok!(Staking::payout_stakers(
+			Origin::signed(1337),
+			ledger.stash,
+			era
+		));
 	}
 }