Session keys buffered for a session. (#2946)

* Session keys buffered for the duration of a session.

* Add queued_keys getter.

* Make sure genesis state is consistent.

* Add validator_count validators.

* Compensate for session delay.

* Remove unused code.

* Add num_validators option.

* Fix session numbers.

* Fix merge.

* Reintroduce changed.

* Update runtime.

* Make NextKeyFor private.

* Move block initialization to function.

* Update lib.rs

* Add test for change propagation.

* Fix docstring.

* Use get instead of take.

* Initialize validators from keys.

* Next try.

* Fix build.

* Fix warning.

* Make initial validator selection more transparent.

* Make storage items private.

* Reorder genesis initialization.

* Update Cargo.lock

* Update runtime version.

* Update runtime version.

* Update Cargo.lock

* Update runtime version.

* Add docs.

substrate/node/cli/src/chain_spec.rs

@@ -112,7 +112,6 @@ fn staging_testnet_config_genesis() -> GenesisConfig {
 				.collect::<Vec<_>>(),
 		}),
 		session: Some(SessionConfig {
-			validators: initial_authorities.iter().map(|x| x.1.clone()).collect(),
 			keys: initial_authorities.iter().map(|x| (x.1.clone(), SessionKeys(x.2.clone(),x.2.clone()))).collect::<Vec<_>>(),
 		}),
 		staking: Some(StakingConfig {
@@ -240,7 +239,6 @@ pub fn testnet_genesis(
 			vesting: vec![],
 		}),
 		session: Some(SessionConfig {
-			validators: initial_authorities.iter().map(|x| x.1.clone()).collect(),
 			keys: initial_authorities.iter().map(|x| (x.1.clone(), SessionKeys(x.2.clone(), x.2.clone()))).collect::<Vec<_>>(),
 		}),
 		staking: Some(StakingConfig {

substrate/node/executor/src/lib.rs

@@ -294,7 +294,6 @@ mod tests {
 	}
 
 	fn new_test_ext(code: &[u8], support_changes_trie: bool) -> TestExternalities<Blake2Hasher> {
-		let three = AccountId::from_raw([3u8; 32]);
 		let mut ext = TestExternalities::new_with_code(code, GenesisConfig {
 			aura: Some(Default::default()),
 			system: Some(SystemConfig {
@@ -319,7 +318,6 @@ mod tests {
 				vesting: vec![],
 			}),
 			session: Some(SessionConfig {
-				validators: vec![AccountKeyring::One.into(), AccountKeyring::Two.into(), three],
 				keys: vec![
 					(alice(), to_session_keys(&AuthorityKeyring::Alice)),
 					(bob(), to_session_keys(&AuthorityKeyring::Bob)),

substrate/node/runtime/src/lib.rs

@@ -71,8 +71,8 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
 	// and set impl_version to equal spec_version. If only runtime
 	// implementation changes and behavior does not, then leave spec_version as
 	// is and increment impl_version.
-	spec_version: 103,
-	impl_version: 103,
+	spec_version: 104,
+	impl_version: 104,
 	apis: RUNTIME_API_VERSIONS,
 };
 
@@ -197,6 +197,7 @@ impl session::Trait for Runtime {
 	type ShouldEndSession = session::PeriodicSessions<Period, Offset>;
 	type Event = Event;
 	type Keys = SessionKeys;
+	type SelectInitialValidators = Staking;
 }
 
 parameter_types! {

substrate/srml/session/Cargo.toml

@@ -10,13 +10,13 @@ safe-mix = { version = "1.0", default-features = false}
 parity-codec = { version = "4.1.1", default-features = false, features = ["derive"] }
 rstd = { package = "sr-std", path = "../../core/sr-std", default-features = false }
 primitives = { package = "sr-primitives", path = "../../core/sr-primitives", default-features = false }
+runtime_io = { package = "sr-io", path = "../../core/sr-io", default-features = false }
 srml-support = { path = "../support", default-features = false }
 system = { package = "srml-system", path = "../system", default-features = false }
 timestamp = { package = "srml-timestamp", path = "../timestamp", default-features = false }
 
 [dev-dependencies]
 substrate-primitives = { path = "../../core/primitives" }
-runtime_io = { package = "sr-io", path = "../../core/sr-io" }
 lazy_static = "1.0"
 
 [features]

substrate/srml/session/src/lib.rs

@@ -118,6 +118,8 @@
 use rstd::{prelude::*, marker::PhantomData, ops::{Sub, Rem}};
 #[cfg(not(feature = "std"))]
 use rstd::alloc::borrow::ToOwned;
+#[cfg(feature = "std")]
+use runtime_io::with_storage;
 use parity_codec::Decode;
 use primitives::traits::{Zero, Member, OpaqueKeys};
 use srml_support::{
@@ -210,6 +212,21 @@ macro_rules! impl_session_handlers {
 
 for_each_tuple!(impl_session_handlers);
 
+/// Handler for selecting the genesis validator set.
+pub trait SelectInitialValidators<T: Trait> {
+	/// Returns the initial validator set. If `None` is returned
+	/// all accounts that have session keys set in the genesis block
+	/// will be validators.
+	fn select_initial_validators() -> Option<Vec<T::AccountId>>;
+}
+
+/// Implementation of `SelectInitialValidators` that does nothing.
+pub struct AllValidators;
+impl<T: Trait> SelectInitialValidators<T> for AllValidators {
+	fn select_initial_validators() -> Option<Vec<T::AccountId>> {
+		None
+	}
+}
 
 pub trait Trait: system::Trait {
 	/// The overarching event type.
@@ -226,6 +243,9 @@ pub trait Trait: system::Trait {
 
 	/// The keys.
 	type Keys: OpaqueKeys + Member + Parameter + Default;
+
+	/// Select initial validators.
+	type SelectInitialValidators: SelectInitialValidators<Self>;
 }
 
 type OpaqueKey = Vec<u8>;
@@ -233,31 +253,64 @@ type OpaqueKey = Vec<u8>;
 decl_storage! {
 	trait Store for Module<T: Trait> as Session {
 		/// The current set of validators.
-		pub Validators get(validators) config(): Vec<T::AccountId>;
+		Validators get(validators): Vec<T::AccountId>;
 
 		/// Current index of the session.
-		pub CurrentIndex get(current_index): SessionIndex;
+		CurrentIndex get(current_index): SessionIndex;
 
 		/// True if anything has changed in this session.
 		Changed: bool;
 
-		/// The next key for a given validator.
-		NextKeyFor build(|config: &GenesisConfig<T>| {
-			config.keys.clone()
-		}): map T::AccountId => Option<T::Keys>;
+		/// The next key to be used for a given validator. At the end of the session, they
+		/// will be moved into the `QueuedKeys` and so changes here will not take effect for
+		/// at least one whole session.
+		NextKeyFor get(next_key_for): map T::AccountId => Option<T::Keys>;
+
+		/// Queued keys changed.
+		QueuedChanged: bool;
+
+		/// The queued keys for the next session. When the next session begins, these keys
+		/// will be used to determine the validator's session keys.
+		QueuedKeys get(queued_keys): Vec<(T::AccountId, T::Keys)>;
 
 		/// The keys that are currently active.
-		Active build(|config: &GenesisConfig<T>| {
-			(0..T::Keys::count()).map(|i| (
-				i as u32,
-				config.keys.iter()
-					.map(|x| x.1.get_raw(i).to_vec())
-					.collect::<Vec<OpaqueKey>>(),
-			)).collect::<Vec<(u32, Vec<OpaqueKey>)>>()
-		}): map u32 => Vec<OpaqueKey>;
+		Active: map u32 => Vec<OpaqueKey>;
 	}
 	add_extra_genesis {
 		config(keys): Vec<(T::AccountId, T::Keys)>;
+		build(|storage, _, config: &GenesisConfig<T>| {
+			with_storage(storage, || {
+				let all_validators = config.keys.iter()
+					.map(|(validator, _)| validator.to_owned())
+					.collect::<Vec<_>>();
+				let all_keys = (0..T::Keys::count()).map(|i| (
+					i as u32,
+					config.keys.iter()
+						.map(|x| x.1.get_raw(i).to_vec())
+						.collect::<Vec<OpaqueKey>>(),
+				)).collect::<Vec<(u32, Vec<OpaqueKey>)>>();
+				<Validators<T>>::put(all_validators.clone());
+				for (v, sk) in config.keys.clone() {
+					<NextKeyFor<T>>::insert(v, sk);
+				}
+				for (i, keys) in all_keys {
+					Active::insert(i, keys);
+				}
+
+				let selected_validators =
+					T::SelectInitialValidators::select_initial_validators()
+					.unwrap_or(all_validators);
+				let selected_keys = selected_validators.iter().map(|validator| {
+					(
+						validator.to_owned(),
+						<NextKeyFor<T>>::get(validator)
+							.unwrap_or_default()
+					)
+				}).collect::<Vec<_>>();
+				<Validators<T>>::put(selected_validators);
+				<QueuedKeys<T>>::put(selected_keys);
+			})
+		})
 	}
 }
 
@@ -338,33 +391,47 @@ decl_module! {
 }
 
 impl<T: Trait> Module<T> {
-	/// Move on to next session: register the new authority set.
+	/// Move on to next session. Register new validator set and session keys. Changes
+	/// to the validator set have a session of delay to take effect. This allows for
+	/// equivocation punishment after a fork.
 	pub fn rotate_session() {
-		// Increment current session index.
 		let session_index = CurrentIndex::get();
 
-		let mut changed = Changed::take();
+		let changed = QueuedChanged::get();
+		let mut next_changed = Changed::take();
 
-		// See if we have a new validator set.
-		let validators = if let Some(new) = T::OnSessionEnding::on_session_ending(session_index) {
-			changed = true;
-			<Validators<T>>::put(&new);
-			new
+		// Get queued session keys and validators.
+		let session_keys = <QueuedKeys<T>>::get();
+		let validators = session_keys.iter()
+			.map(|(validator, _)| validator.to_owned())
+			.collect::<Vec<_>>();
+		<Validators<T>>::put(validators);
+
+		// Get next validator set.
+		let maybe_validators = T::OnSessionEnding::on_session_ending(session_index);
+		let next_validators = if let Some(validators) = maybe_validators {
+			next_changed = true;
+			validators
 		} else {
 			<Validators<T>>::get()
 		};
 
+		// Increment session index.
 		let session_index = session_index + 1;
 		CurrentIndex::put(session_index);
 
+		// Queue next session keys.
+		let next_session_keys = next_validators.into_iter()
+			.map(|a| { let k = <NextKeyFor<T>>::get(&a).unwrap_or_default(); (a, k) })
+			.collect::<Vec<_>>();
+		<QueuedKeys<T>>::put(next_session_keys);
+		QueuedChanged::put(next_changed);
+
 		// Record that this happened.
 		Self::deposit_event(Event::NewSession(session_index));
 
 		// Tell everyone about the new session keys.
-		let amalgamated = validators.into_iter()
-			.map(|a| { let k = <NextKeyFor<T>>::get(&a).unwrap_or_default(); (a, k) })
-			.collect::<Vec<_>>();
-		T::SessionHandler::on_new_session::<T::Keys>(changed, &amalgamated);
+		T::SessionHandler::on_new_session::<T::Keys>(changed, &session_keys);
 	}
 
 	/// Disable the validator of index `i`.
@@ -425,6 +492,8 @@ mod tests {
 			RefCell::new(vec![UintAuthorityId(1), UintAuthorityId(2), UintAuthorityId(3)]);
 		static FORCE_SESSION_END: RefCell<bool> = RefCell::new(false);
 		static SESSION_LENGTH: RefCell<u64> = RefCell::new(2);
+		static SESSION_CHANGED: RefCell<bool> = RefCell::new(false);
+		static TEST_SESSION_CHANGED: RefCell<bool> = RefCell::new(false);
 	}
 
 	pub struct TestShouldEndSession;
@@ -437,7 +506,8 @@ mod tests {
 
 	pub struct TestSessionHandler;
 	impl SessionHandler<u64> for TestSessionHandler {
-		fn on_new_session<T: OpaqueKeys>(_changed: bool, validators: &[(u64, T)]) {
+		fn on_new_session<T: OpaqueKeys>(changed: bool, validators: &[(u64, T)]) {
+			SESSION_CHANGED.with(|l| *l.borrow_mut() = changed);
 			AUTHORITIES.with(|l|
 				*l.borrow_mut() = validators.iter().map(|(_, id)| id.get::<UintAuthorityId>(0).unwrap_or_default()).collect()
 			);
@@ -448,7 +518,11 @@ mod tests {
 	pub struct TestOnSessionEnding;
 	impl OnSessionEnding<u64> for TestOnSessionEnding {
 		fn on_session_ending(_: SessionIndex) -> Option<Vec<u64>> {
-			Some(NEXT_VALIDATORS.with(|l| l.borrow().clone()))
+			if !TEST_SESSION_CHANGED.with(|l| *l.borrow()) {
+				Some(NEXT_VALIDATORS.with(|l| l.borrow().clone()))
+			} else {
+				None
+			}
 		}
 	}
 
@@ -464,6 +538,10 @@ mod tests {
 		SESSION_LENGTH.with(|l| *l.borrow_mut() = x )
 	}
 
+	fn session_changed() -> bool {
+		SESSION_CHANGED.with(|l| *l.borrow())
+	}
+
 	#[derive(Clone, Eq, PartialEq)]
 	pub struct Test;
 	impl system::Trait for Test {
@@ -487,18 +565,19 @@ mod tests {
 		type SessionHandler = TestSessionHandler;
 		type Keys = UintAuthorityId;
 		type Event = ();
+		type SelectInitialValidators = AllValidators;
 	}
 
 	type System = system::Module<Test>;
 	type Session = Module<Test>;
 
 	fn new_test_ext() -> runtime_io::TestExternalities<Blake2Hasher> {
+		TEST_SESSION_CHANGED.with(|l| *l.borrow_mut() = false);
 		let mut t = system::GenesisConfig::default().build_storage::<Test>().unwrap().0;
 		t.extend(timestamp::GenesisConfig::<Test> {
 			minimum_period: 5,
 		}.build_storage().unwrap().0);
 		t.extend(GenesisConfig::<Test> {
-			validators: NEXT_VALIDATORS.with(|l| l.borrow().clone()),
 			keys: NEXT_VALIDATORS.with(|l|
 				l.borrow().iter().cloned().map(|i| (i, UintAuthorityId(i))).collect()
 			),
@@ -506,6 +585,12 @@ mod tests {
 		runtime_io::TestExternalities::new(t)
 	}
 
+	fn initialize_block(block: u64) {
+		SESSION_CHANGED.with(|l| *l.borrow_mut() = false);
+		System::set_block_number(block);
+		Session::on_initialize(block);
+	}
+
 	#[test]
 	fn simple_setup_should_work() {
 		with_externalities(&mut new_test_ext(), || {
@@ -518,28 +603,46 @@ mod tests {
 	fn authorities_should_track_validators() {
 		with_externalities(&mut new_test_ext(), || {
 			NEXT_VALIDATORS.with(|v| *v.borrow_mut() = vec![1, 2]);
-			force_new_session();
 
-			System::set_block_number(1);
-			Session::on_initialize(1);
+			force_new_session();
+			initialize_block(1);
+			assert_eq!(Session::queued_keys(), vec![
+				(1, UintAuthorityId(1)),
+				(2, UintAuthorityId(2)),
+			]);
+			assert_eq!(Session::validators(), vec![1, 2, 3]);
+			assert_eq!(authorities(), vec![UintAuthorityId(1), UintAuthorityId(2), UintAuthorityId(3)]);
+
+			force_new_session();
+			initialize_block(2);
+			assert_eq!(Session::queued_keys(), vec![
+				(1, UintAuthorityId(1)),
+				(2, UintAuthorityId(2)),
+			]);
 			assert_eq!(Session::validators(), vec![1, 2]);
 			assert_eq!(authorities(), vec![UintAuthorityId(1), UintAuthorityId(2)]);
 
 			NEXT_VALIDATORS.with(|v| *v.borrow_mut() = vec![1, 2, 4]);
 			assert_ok!(Session::set_keys(Origin::signed(4), UintAuthorityId(4), vec![]));
+			force_new_session();
+			initialize_block(3);
+			assert_eq!(Session::queued_keys(), vec![
+				(1, UintAuthorityId(1)),
+				(2, UintAuthorityId(2)),
+				(4, UintAuthorityId(4)),
+			]);
+			assert_eq!(Session::validators(), vec![1, 2]);
+			assert_eq!(authorities(), vec![UintAuthorityId(1), UintAuthorityId(2)]);
 
 			force_new_session();
-			System::set_block_number(2);
-			Session::on_initialize(2);
+			initialize_block(4);
+			assert_eq!(Session::queued_keys(), vec![
+				(1, UintAuthorityId(1)),
+				(2, UintAuthorityId(2)),
+				(4, UintAuthorityId(4)),
+			]);
 			assert_eq!(Session::validators(), vec![1, 2, 4]);
 			assert_eq!(authorities(), vec![UintAuthorityId(1), UintAuthorityId(2), UintAuthorityId(4)]);
-
-			NEXT_VALIDATORS.with(|v| *v.borrow_mut() = vec![1, 2, 3]);
-			force_new_session();
-			System::set_block_number(3);
-			Session::on_initialize(3);
-			assert_eq!(Session::validators(), vec![1, 2, 3]);
-			assert_eq!(authorities(), vec![UintAuthorityId(1), UintAuthorityId(2), UintAuthorityId(3)]);
 		});
 	}
 
@@ -548,25 +651,20 @@ mod tests {
 		with_externalities(&mut new_test_ext(), || {
 			set_session_length(10);
 
-			System::set_block_number(1);
-			Session::on_initialize(1);
+			initialize_block(1);
 			assert_eq!(Session::current_index(), 0);
 
-			System::set_block_number(2);
-			Session::on_initialize(2);
+			initialize_block(2);
 			assert_eq!(Session::current_index(), 0);
+
 			force_new_session();
-
-			System::set_block_number(3);
-			Session::on_initialize(3);
+			initialize_block(3);
 			assert_eq!(Session::current_index(), 1);
 
-			System::set_block_number(9);
-			Session::on_initialize(9);
+			initialize_block(9);
 			assert_eq!(Session::current_index(), 1);
 
-			System::set_block_number(10);
-			Session::on_initialize(10);
+			initialize_block(10);
 			assert_eq!(Session::current_index(), 2);
 		});
 	}
@@ -575,28 +673,69 @@ mod tests {
 	fn session_change_should_work() {
 		with_externalities(&mut new_test_ext(), || {
 			// Block 1: No change
-			System::set_block_number(1);
-			Session::on_initialize(1);
+			initialize_block(1);
 			assert_eq!(authorities(), vec![UintAuthorityId(1), UintAuthorityId(2), UintAuthorityId(3)]);
 
 			// Block 2: Session rollover, but no change.
-			System::set_block_number(2);
-			Session::on_initialize(2);
+			initialize_block(2);
 			assert_eq!(authorities(), vec![UintAuthorityId(1), UintAuthorityId(2), UintAuthorityId(3)]);
 
 			// Block 3: Set new key for validator 2; no visible change.
-			System::set_block_number(3);
-			Session::on_initialize(3);
+			initialize_block(3);
 			assert_ok!(Session::set_keys(Origin::signed(2), UintAuthorityId(5), vec![]));
 			assert_eq!(authorities(), vec![UintAuthorityId(1), UintAuthorityId(2), UintAuthorityId(3)]);
 
-			// Block 4: Session rollover, authority 2 changes.
-			System::set_block_number(4);
-			Session::on_initialize(4);
+			// Block 4: Session rollover; no visible change.
+			initialize_block(4);
+			assert_eq!(authorities(), vec![UintAuthorityId(1), UintAuthorityId(2), UintAuthorityId(3)]);
+
+			// Block 5: No change.
+			initialize_block(5);
+			assert_eq!(authorities(), vec![UintAuthorityId(1), UintAuthorityId(2), UintAuthorityId(3)]);
+
+			// Block 6: Session rollover; authority 2 changes.
+			initialize_block(6);
 			assert_eq!(authorities(), vec![UintAuthorityId(1), UintAuthorityId(5), UintAuthorityId(3)]);
 		});
 	}
 
+	#[test]
+	fn session_changed_flag_works() {
+		with_externalities(&mut new_test_ext(), || {
+			TEST_SESSION_CHANGED.with(|l| *l.borrow_mut() = true);
+
+			force_new_session();
+			initialize_block(1);
+			assert!(!session_changed());
+
+			force_new_session();
+			initialize_block(2);
+			assert!(!session_changed());
+
+			Session::disable_index(0);
+			force_new_session();
+			initialize_block(3);
+			assert!(!session_changed());
+
+			force_new_session();
+			initialize_block(4);
+			assert!(session_changed());
+
+			force_new_session();
+			initialize_block(5);
+			assert!(!session_changed());
+
+			assert_ok!(Session::set_keys(Origin::signed(2), UintAuthorityId(5), vec![]));
+			force_new_session();
+			initialize_block(6);
+			assert!(!session_changed());
+
+			force_new_session();
+			initialize_block(7);
+			assert!(session_changed());
+		});
+	}
+
 	#[test]
 	fn periodic_session_works() {
 		struct Period;

substrate/srml/staking/src/lib.rs

@@ -281,7 +281,7 @@ use srml_support::{
 		WithdrawReasons, OnUnbalanced, Imbalance, Get
 	}
 };
-use session::{OnSessionEnding, SessionIndex};
+use session::{OnSessionEnding, SelectInitialValidators, SessionIndex};
 use primitives::Perbill;
 use primitives::traits::{
 	Convert, Zero, One, StaticLookup, CheckedSub, CheckedShl, Saturating, Bounded
@@ -587,10 +587,6 @@ decl_storage! {
 						}, _ => Ok(())
 					};
 				}
-
-				if let (_, Some(validators)) = <Module<T>>::select_validators() {
-					<session::Validators<T>>::put(&validators);
-				}
 			});
 		});
 	}
@@ -1255,3 +1251,9 @@ impl<T: Trait> OnFreeBalanceZero<T::AccountId> for Module<T> {
 		<Nominators<T>>::remove(stash);
 	}
 }
+
+impl<T: Trait> SelectInitialValidators<T> for Module<T> {
+	fn select_initial_validators() -> Option<Vec<T::AccountId>> {
+		<Module<T>>::select_validators().1
+	}
+}

substrate/srml/staking/src/mock.rs

@@ -126,6 +126,7 @@ impl session::Trait for Test {
 	type ShouldEndSession = session::PeriodicSessions<Period, Offset>;
 	type SessionHandler = TestSessionHandler;
 	type Event = ();
+	type SelectInitialValidators = Staking;
 }
 impl timestamp::Trait for Test {
 	type Moment = u64;
@@ -148,26 +149,26 @@ impl Trait for Test {
 
 pub struct ExtBuilder {
 	existential_deposit: u64,
-	current_era: EraIndex,
 	reward: u64,
 	validator_pool: bool,
 	nominate: bool,
 	validator_count: u32,
 	minimum_validator_count: u32,
 	fair: bool,
+	num_validators: Option<u32>,
 }
 
 impl Default for ExtBuilder {
 	fn default() -> Self {
 		Self {
 			existential_deposit: 0,
-			current_era: 0,
 			reward: 10,
 			validator_pool: false,
 			nominate: true,
 			validator_count: 2,
 			minimum_validator_count: 0,
-			fair: true
+			fair: true,
+			num_validators: None,
 		}
 	}
 }
@@ -177,10 +178,6 @@ impl ExtBuilder {
 		self.existential_deposit = existential_deposit;
 		self
 	}
-	pub fn _current_era(mut self, current_era: EraIndex) -> Self {
-		self.current_era = current_era;
-		self
-	}
 	pub fn validator_pool(mut self, validator_pool: bool) -> Self {
 		self.validator_pool = validator_pool;
 		self
@@ -201,6 +198,10 @@ impl ExtBuilder {
 		self.fair = is_fair;
 		self
 	}
+	pub fn num_validators(mut self, num_validators: u32) -> Self {
+		self.num_validators = Some(num_validators);
+		self
+	}
 	pub fn set_associated_consts(&self) {
 		EXISTENTIAL_DEPOSIT.with(|v| *v.borrow_mut() = self.existential_deposit);
 	}
@@ -212,12 +213,12 @@ impl ExtBuilder {
 		} else {
 			1
 		};
-		let validators = if self.validator_pool { vec![10, 20, 30, 40] } else { vec![10, 20] };
-		let _ = session::GenesisConfig::<Test>{
-			// NOTE: if config.nominate == false then 100 is also selected in the initial round.
-			validators,
-			keys: vec![],
-		}.assimilate_storage(&mut t, &mut c);
+
+		let num_validators = self.num_validators.unwrap_or(self.validator_count);
+		let validators = (0..num_validators)
+			.map(|x| ((x + 1) * 10) as u64)
+			.collect::<Vec<_>>();
+
 		let _ = balances::GenesisConfig::<Test>{
 			balances: vec![
 					(1, 10 * balance_factor),
@@ -237,6 +238,7 @@ impl ExtBuilder {
 			],
 			vesting: vec![],
 		}.assimilate_storage(&mut t, &mut c);
+
 		let stake_21 = if self.fair { 1000 } else { 2000 };
 		let stake_31 = if self.validator_pool { balance_factor * 1000 } else { 1 };
 		let status_41 = if self.validator_pool {
@@ -246,7 +248,7 @@ impl ExtBuilder {
 		};
 		let nominated = if self.nominate { vec![11, 21] } else { vec![] };
 		let _ = GenesisConfig::<Test>{
-			current_era: self.current_era,
+			current_era: 0,
 			stakers: vec![
 				(11, 10, balance_factor * 1000, StakerStatus::<AccountId>::Validator),
 				(21, 20, stake_21, StakerStatus::<AccountId>::Validator),
@@ -263,9 +265,15 @@ impl ExtBuilder {
 			offline_slash_grace: 0,
 			invulnerables: vec![],
 		}.assimilate_storage(&mut t, &mut c);
+
 		let _ = timestamp::GenesisConfig::<Test>{
 			minimum_period: 5,
 		}.assimilate_storage(&mut t, &mut c);
+
+		let _ = session::GenesisConfig::<Test> {
+			keys: validators.iter().map(|x| (*x, UintAuthorityId(*x))).collect(),
+		}.assimilate_storage(&mut t, &mut c);
+
 		let mut ext = t.into();
 		runtime_io::with_externalities(&mut ext, || {
 			let validators = Session::validators();
@@ -346,6 +354,8 @@ pub fn bond_nominator(acc: u64, val: u64, target: Vec<u64>) {
 }
 
 pub fn start_session(session_index: session::SessionIndex) {
+	// Compensate for session delay
+	let session_index = session_index + 1;
 	for i in 0..(session_index - Session::current_index()) {
 		System::set_block_number((i + 1).into());
 		Session::on_initialize(System::block_number());

substrate/srml/staking/src/tests.rs

@@ -411,18 +411,18 @@ fn multi_era_reward_should_work() {
 		// Set payee to controller
 		assert_ok!(Staking::set_payee(Origin::signed(10), RewardDestination::Controller));
 
-		start_session(1);
+		start_session(0);
 
 		// session triggered: the reward value stashed should be 10
 		assert_eq!(Staking::current_session_reward(), session_reward);
 		assert_eq!(Staking::current_era_reward(), session_reward);
 
-		start_session(2);
+		start_session(1);
 
 		assert_eq!(Staking::current_session_reward(), session_reward);
 		assert_eq!(Staking::current_era_reward(), 2*session_reward);
 
-		start_session(3);
+		start_session(2);
 
 		// 1 + sum of of the session rewards accumulated
 		let recorded_balance = 1 + 3*session_reward;
@@ -433,13 +433,13 @@ fn multi_era_reward_should_work() {
 		assert_eq!(Staking::current_session_reward(), new_session_reward);
 
 		// fast forward to next era:
-		start_session(5);
+		start_session(4);
 
 		// intermediate test.
 		assert_eq!(Staking::current_era_reward(), 2*new_session_reward);
 
 		// new era is triggered here.
-		start_session(6);
+		start_session(5);
 
 		// pay time
 		assert_eq!(Balances::total_balance(&10), 3*new_session_reward + recorded_balance);
@@ -464,10 +464,7 @@ fn staking_should_work() {
 		for i in 1..5 { let _ = Balances::make_free_balance_be(&i, 2000); }
 
 		// --- Block 1:
-		System::set_block_number(1);
-		Session::on_initialize(System::block_number());
-		assert_eq!(Staking::current_era(), 0);
-
+		start_session(1);
 		// add a new candidate for being a validator. account 3 controlled by 4.
 		assert_ok!(Staking::bond(Origin::signed(3), 4, 1500, RewardDestination::Controller));
 		assert_ok!(Staking::validate(Origin::signed(4), ValidatorPrefs::default()));
@@ -476,48 +473,30 @@ fn staking_should_work() {
 		assert_eq_uvec!(Session::validators(), vec![20, 10]);
 
 		// --- Block 2:
-		System::set_block_number(2);
-		Session::on_initialize(System::block_number());
-		assert_eq!(Staking::current_era(), 0);
+		start_session(2);
 
 		// No effects will be seen so far. Era has not been yet triggered.
 		assert_eq_uvec!(Session::validators(), vec![20, 10]);
 
 
-		// --- Block 3: the validators will now change.
-		System::set_block_number(3);
-		Session::on_initialize(System::block_number());
-
-		// 2 only voted for 4 and 20
-		assert_eq!(Session::validators().len(), 2);
-		assert_eq_uvec!(Session::validators(), vec![20, 4]);
+		// --- Block 3: the validators will now be queued.
+		start_session(3);
 		assert_eq!(Staking::current_era(), 1);
 
+		// --- Block 4: the validators will now be changed.
+		start_session(4);
 
+		assert_eq_uvec!(Session::validators(), vec![20, 4]);
 		// --- Block 4: Unstake 4 as a validator, freeing up the balance stashed in 3
-		System::set_block_number(4);
-		Session::on_initialize(System::block_number());
-
 		// 4 will chill
 		Staking::chill(Origin::signed(4)).unwrap();
 
-		// nothing should be changed so far.
-		assert_eq_uvec!(Session::validators(), vec![20, 4]);
-		assert_eq!(Staking::current_era(), 1);
-
-
 		// --- Block 5: nothing. 4 is still there.
-		System::set_block_number(5);
-		Session::on_initialize(System::block_number());
+		start_session(5);
 		assert_eq_uvec!(Session::validators(), vec![20, 4]);
-		assert_eq!(Staking::current_era(), 1);
-
 
 		// --- Block 6: 4 will not be a validator.
-		System::set_block_number(6);
-		Session::on_initialize(System::block_number());
-		assert_eq!(Staking::current_era(), 2);
-		assert_eq!(Session::validators().contains(&4), false);
+		start_session(7);
 		assert_eq_uvec!(Session::validators(), vec![20, 10]);
 
 		// Note: the stashed value of 4 is still lock
@@ -537,6 +516,7 @@ fn less_than_needed_candidates_works() {
 		.minimum_validator_count(1)
 		.validator_count(4)
 		.nominate(false)
+		.num_validators(3)
 		.build(),
 	|| {
 		assert_eq!(Staking::validator_count(), 4);
@@ -564,6 +544,7 @@ fn no_candidate_emergency_condition() {
 	with_externalities(&mut ExtBuilder::default()
 		.minimum_validator_count(10)
 		.validator_count(15)
+		.num_validators(4)
 		.validator_pool(true)
 		.nominate(false)
 		.build(),
@@ -857,37 +838,37 @@ fn session_and_eras_work() {
 		assert_eq!(Staking::current_era(), 0);
 
 		// Block 1: No change.
-		start_session(1);
+		start_session(0);
 		assert_eq!(Session::current_index(), 1);
 		assert_eq!(Staking::current_era(), 0);
 
 		// Block 2: Simple era change.
-		start_session(3);
+		start_session(2);
 		assert_eq!(Session::current_index(), 3);
 		assert_eq!(Staking::current_era(), 1);
 
 		// Block 3: Schedule an era length change; no visible changes.
-		start_session(4);
+		start_session(3);
 		assert_eq!(Session::current_index(), 4);
 		assert_eq!(Staking::current_era(), 1);
 
 		// Block 4: Era change kicks in.
-		start_session(6);
+		start_session(5);
 		assert_eq!(Session::current_index(), 6);
 		assert_eq!(Staking::current_era(), 2);
 
 		// Block 5: No change.
-		start_session(7);
+		start_session(6);
 		assert_eq!(Session::current_index(), 7);
 		assert_eq!(Staking::current_era(), 2);
 
 		// Block 6: No change.
-		start_session(8);
+		start_session(7);
 		assert_eq!(Session::current_index(), 8);
 		assert_eq!(Staking::current_era(), 2);
 
 		// Block 7: Era increment.
-		start_session(9);
+		start_session(8);
 		assert_eq!(Session::current_index(), 9);
 		assert_eq!(Staking::current_era(), 3);
 	});
@@ -1656,22 +1637,19 @@ fn switching_roles() {
 		assert_ok!(Staking::validate(Origin::signed(6), ValidatorPrefs::default()));
 
 		// new block
-		System::set_block_number(1);
-		Session::on_initialize(System::block_number());
+		start_session(1);
 
 		// no change
 		assert_eq_uvec!(Session::validators(), vec![20, 10]);
 
 		// new block
-		System::set_block_number(2);
-		Session::on_initialize(System::block_number());
+		start_session(2);
 
 		// no change
 		assert_eq_uvec!(Session::validators(), vec![20, 10]);
 
 		// new block --> ne era --> new validators
-		System::set_block_number(3);
-		Session::on_initialize(System::block_number());
+		start_session(3);
 
 		// with current nominators 10 and 5 have the most stake
 		assert_eq_uvec!(Session::validators(), vec![6, 10]);
@@ -1685,17 +1663,14 @@ fn switching_roles() {
 		// 2 : 2000 self vote + 250 vote.
 		// Winners: 20 and 2
 
-		System::set_block_number(4);
-		Session::on_initialize(System::block_number());
+		start_session(4);
 		assert_eq_uvec!(Session::validators(), vec![6, 10]);
 
-		System::set_block_number(5);
-		Session::on_initialize(System::block_number());
+		start_session(5);
 		assert_eq_uvec!(Session::validators(), vec![6, 10]);
 
 		// ne era
-		System::set_block_number(6);
-		Session::on_initialize(System::block_number());
+		start_session(6);
 		assert_eq_uvec!(Session::validators(), vec![2, 20]);
 
 		check_exposure_all();