Improves the EPM/Staking e2e test setup (#14292)

* Adds onchainify to e2e tests; Adds try-runtime checks at the end of the tests

* Refactors OCW to be more realistic and to submit solutions during unsigned phase

* Uses config palllet Default impl in mock

* Improves OCW progression vode

* simplified OCW tests

substrate/Cargo.lock

@@ -6503,6 +6503,7 @@ dependencies = [
  "pallet-staking",
  "pallet-timestamp",
  "parity-scale-codec",
+ "parking_lot 0.12.1",
  "scale-info",
  "sp-core",
  "sp-io",

substrate/frame/election-provider-multi-phase/test-staking-e2e/Cargo.toml

@@ -13,6 +13,7 @@ publish = false
 targets = ["x86_64-unknown-linux-gnu"]
 
 [dev-dependencies]
+parking_lot = "0.12.1"
 codec = { package = "parity-scale-codec", version = "3.6.1", features = ["derive"] }
 scale-info = { version = "2.0.1", features = ["derive"] }
 log = { version = "0.4.17", default-features = false }

substrate/frame/election-provider-multi-phase/test-staking-e2e/src/lib.rs

@@ -53,12 +53,14 @@ fn log_current_time() {
 
 #[test]
 fn block_progression_works() {
-	ExtBuilder::default().build_and_execute(|| {
+	let (mut ext, pool_state, _) = ExtBuilder::default().build_offchainify();
+
+	ext.execute_with(|| {
 		assert_eq!(active_era(), 0);
 		assert_eq!(Session::current_index(), 0);
 		assert!(ElectionProviderMultiPhase::current_phase().is_off());
 
-		assert!(start_next_active_era().is_ok());
+		assert!(start_next_active_era(pool_state.clone()).is_ok());
 		assert_eq!(active_era(), 1);
 		assert_eq!(Session::current_index(), <SessionsPerEra as Get<u32>>::get());
 
@@ -68,12 +70,14 @@ fn block_progression_works() {
 		assert!(ElectionProviderMultiPhase::current_phase().is_signed());
 	});
 
-	ExtBuilder::default().build_and_execute(|| {
+	let (mut ext, pool_state, _) = ExtBuilder::default().build_offchainify();
+
+	ext.execute_with(|| {
 		assert_eq!(active_era(), 0);
 		assert_eq!(Session::current_index(), 0);
 		assert!(ElectionProviderMultiPhase::current_phase().is_off());
 
-		assert!(start_next_active_era_delayed_solution().is_ok());
+		assert!(start_next_active_era_delayed_solution(pool_state).is_ok());
 		// if the solution is delayed, EPM will end up in emergency mode..
 		assert!(ElectionProviderMultiPhase::current_phase().is_emergency());
 		// .. era won't progress..
@@ -83,6 +87,45 @@ fn block_progression_works() {
 	})
 }
 
+#[test]
+fn offchainify_works() {
+	use pallet_election_provider_multi_phase::QueuedSolution;
+
+	let staking_builder = StakingExtBuilder::default();
+	let epm_builder = EpmExtBuilder::default();
+	let (mut ext, pool_state, _) = ExtBuilder::default()
+		.epm(epm_builder)
+		.staking(staking_builder)
+		.build_offchainify();
+
+	ext.execute_with(|| {
+		// test ocw progression and solution queue if submission when unsigned phase submission is
+		// not delayed.
+		for _ in 0..100 {
+			roll_one(pool_state.clone(), false);
+			let current_phase = ElectionProviderMultiPhase::current_phase();
+
+			assert!(
+				match QueuedSolution::<Runtime>::get() {
+					Some(_) => current_phase.is_unsigned(),
+					None => !current_phase.is_unsigned(),
+				},
+				"solution must be queued *only* in unsigned phase"
+			);
+		}
+
+		// test ocw solution queue if submission in unsigned phase is delayed.
+		for _ in 0..100 {
+			roll_one(pool_state.clone(), true);
+			assert_eq!(
+				QueuedSolution::<Runtime>::get(),
+				None,
+				"solution must never be submitted and stored since it is delayed"
+			);
+		}
+	})
+}
+
 #[test]
 /// Replicates the Kusama incident of 8th Dec 2022 and its resolution through the governance
 /// fallback.
@@ -99,8 +142,9 @@ fn block_progression_works() {
 /// restarts. Note that in this test case, the emergency throttling is disabled.
 fn enters_emergency_phase_after_forcing_before_elect() {
 	let epm_builder = EpmExtBuilder::default().disable_emergency_throttling();
+	let (mut ext, pool_state, _) = ExtBuilder::default().epm(epm_builder).build_offchainify();
 
-	ExtBuilder::default().epm(epm_builder).build_and_execute(|| {
+	ext.execute_with(|| {
 		log!(
 			trace,
 			"current validators (staking): {:?}",
@@ -117,15 +161,15 @@ fn enters_emergency_phase_after_forcing_before_elect() {
 
 		assert_eq!(pallet_staking::ForceEra::<Runtime>::get(), pallet_staking::Forcing::ForceNew);
 
-		advance_session_delayed_solution();
+		advance_session_delayed_solution(pool_state.clone());
 		assert!(ElectionProviderMultiPhase::current_phase().is_emergency());
 		log_current_time();
 
 		let era_before_delayed_next = Staking::current_era();
 		// try to advance 2 eras.
-		assert!(start_next_active_era_delayed_solution().is_ok());
+		assert!(start_next_active_era_delayed_solution(pool_state.clone()).is_ok());
 		assert_eq!(Staking::current_era(), era_before_delayed_next);
-		assert!(start_next_active_era().is_err());
+		assert!(start_next_active_era(pool_state).is_err());
 		assert_eq!(Staking::current_era(), era_before_delayed_next);
 
 		// EPM is still in emergency phase.
@@ -169,41 +213,43 @@ fn continous_slashes_below_offending_threshold() {
 	let staking_builder = StakingExtBuilder::default().validator_count(10);
 	let epm_builder = EpmExtBuilder::default().disable_emergency_throttling();
 
-	ExtBuilder::default()
-		.staking(staking_builder)
+	let (mut ext, pool_state, _) = ExtBuilder::default()
 		.epm(epm_builder)
-		.build_and_execute(|| {
-			assert_eq!(Session::validators().len(), 10);
-			let mut active_validator_set = Session::validators();
+		.staking(staking_builder)
+		.build_offchainify();
 
-			roll_to_epm_signed();
+	ext.execute_with(|| {
+		assert_eq!(Session::validators().len(), 10);
+		let mut active_validator_set = Session::validators();
 
-			// set a minimum election score.
-			assert!(set_minimum_election_score(500, 1000, 500).is_ok());
+		roll_to_epm_signed();
 
-			// slash 10% of the active validators and progress era until the minimum trusted score
-			// is reached.
-			while active_validator_set.len() > 0 {
-				let slashed = slash_percentage(Perbill::from_percent(10));
-				assert_eq!(slashed.len(), 1);
+		// set a minimum election score.
+		assert!(set_minimum_election_score(500, 1000, 500).is_ok());
 
-				// break loop when era does not progress; EPM is in emergency phase as election
-				// failed due to election minimum score.
-				if start_next_active_era().is_err() {
-					assert!(ElectionProviderMultiPhase::current_phase().is_emergency());
-					break
-				}
+		// slash 10% of the active validators and progress era until the minimum trusted score
+		// is reached.
+		while active_validator_set.len() > 0 {
+			let slashed = slash_percentage(Perbill::from_percent(10));
+			assert_eq!(slashed.len(), 1);
 
-				active_validator_set = Session::validators();
-
-				log!(
-					trace,
-					"slashed 10% of active validators ({:?}). After slash: {:?}",
-					slashed,
-					active_validator_set
-				);
+			// break loop when era does not progress; EPM is in emergency phase as election
+			// failed due to election minimum score.
+			if start_next_active_era(pool_state.clone()).is_err() {
+				assert!(ElectionProviderMultiPhase::current_phase().is_emergency());
+				break
 			}
-		});
+
+			active_validator_set = Session::validators();
+
+			log!(
+				trace,
+				"slashed 10% of active validators ({:?}). After slash: {:?}",
+				slashed,
+				active_validator_set
+			);
+		}
+	});
 }
 
 #[test]
@@ -223,54 +269,53 @@ fn set_validation_intention_after_chilled() {
 	use frame_election_provider_support::SortedListProvider;
 	use pallet_staking::{Event, Forcing, Nominators};
 
-	let staking_builder = StakingExtBuilder::default();
-	let epm_builder = EpmExtBuilder::default();
+	let (mut ext, pool_state, _) = ExtBuilder::default()
+		.epm(EpmExtBuilder::default())
+		.staking(StakingExtBuilder::default())
+		.build_offchainify();
 
-	ExtBuilder::default()
-		.staking(staking_builder)
-		.epm(epm_builder)
-		.build_and_execute(|| {
-			assert_eq!(active_era(), 0);
-			// validator is part of the validator set.
-			assert!(Session::validators().contains(&81));
-			assert!(<Runtime as pallet_staking::Config>::VoterList::contains(&81));
+	ext.execute_with(|| {
+		assert_eq!(active_era(), 0);
+		// validator is part of the validator set.
+		assert!(Session::validators().contains(&41));
+		assert!(<Runtime as pallet_staking::Config>::VoterList::contains(&41));
 
-			// nominate validator 81.
-			assert_ok!(Staking::nominate(RuntimeOrigin::signed(21), vec![81]));
-			assert_eq!(Nominators::<Runtime>::get(21).unwrap().targets, vec![81]);
+		// nominate validator 81.
+		assert_ok!(Staking::nominate(RuntimeOrigin::signed(21), vec![41]));
+		assert_eq!(Nominators::<Runtime>::get(21).unwrap().targets, vec![41]);
 
-			// validator is slashed. it is removed from the `VoterList` through chilling but in the
-			// current era, the validator is still part of the active validator set.
-			add_slash(&81);
-			assert!(Session::validators().contains(&81));
-			assert!(!<Runtime as pallet_staking::Config>::VoterList::contains(&81));
-			assert_eq!(
-				staking_events(),
-				[
-					Event::Chilled { stash: 81 },
-					Event::ForceEra { mode: Forcing::ForceNew },
-					Event::SlashReported {
-						validator: 81,
-						slash_era: 0,
-						fraction: Perbill::from_percent(10)
-					}
-				],
-			);
+		// validator is slashed. it is removed from the `VoterList` through chilling but in the
+		// current era, the validator is still part of the active validator set.
+		add_slash(&41);
+		assert!(Session::validators().contains(&41));
+		assert!(!<Runtime as pallet_staking::Config>::VoterList::contains(&41));
+		assert_eq!(
+			staking_events(),
+			[
+				Event::Chilled { stash: 41 },
+				Event::ForceEra { mode: Forcing::ForceNew },
+				Event::SlashReported {
+					validator: 41,
+					slash_era: 0,
+					fraction: Perbill::from_percent(10)
+				}
+			],
+		);
 
-			// after the nominator is slashed and chilled, the nominations remain.
-			assert_eq!(Nominators::<Runtime>::get(21).unwrap().targets, vec![81]);
+		// after the nominator is slashed and chilled, the nominations remain.
+		assert_eq!(Nominators::<Runtime>::get(21).unwrap().targets, vec![41]);
 
-			// validator sets intention to stake again in the same era it was chilled.
-			assert_ok!(Staking::validate(RuntimeOrigin::signed(81), Default::default()));
+		// validator sets intention to stake again in the same era it was chilled.
+		assert_ok!(Staking::validate(RuntimeOrigin::signed(41), Default::default()));
 
-			// progress era and check that the slashed validator is still part of the validator
-			// set.
-			assert!(start_next_active_era().is_ok());
-			assert_eq!(active_era(), 1);
-			assert!(Session::validators().contains(&81));
-			assert!(<Runtime as pallet_staking::Config>::VoterList::contains(&81));
+		// progress era and check that the slashed validator is still part of the validator
+		// set.
+		assert!(start_next_active_era(pool_state).is_ok());
+		assert_eq!(active_era(), 1);
+		assert!(Session::validators().contains(&41));
+		assert!(<Runtime as pallet_staking::Config>::VoterList::contains(&41));
 
-			// nominations are still active as before the slash.
-			assert_eq!(Nominators::<Runtime>::get(21).unwrap().targets, vec![81]);
-		})
+		// nominations are still active as before the slash.
+		assert_eq!(Nominators::<Runtime>::get(21).unwrap().targets, vec![41]);
+	})
 }

substrate/frame/election-provider-multi-phase/test-staking-e2e/src/mock.rs

@@ -18,13 +18,19 @@
 #![allow(dead_code)]
 
 use _feps::ExtendedBalance;
-use frame_support::{parameter_types, traits, traits::Hooks, weights::constants};
+use frame_support::{
+	dispatch::UnfilteredDispatchable, parameter_types, traits, traits::Hooks, weights::constants,
+};
 use frame_system::EnsureRoot;
-use sp_core::{ConstU32, Get, H256};
+use sp_core::{ConstU32, Get};
 use sp_npos_elections::{ElectionScore, VoteWeight};
 use sp_runtime::{
+	offchain::{
+		testing::{OffchainState, PoolState, TestOffchainExt, TestTransactionPoolExt},
+		OffchainDbExt, OffchainWorkerExt, TransactionPoolExt,
+	},
 	testing,
-	traits::{IdentityLookup, Zero},
+	traits::Zero,
 	transaction_validity, BuildStorage, PerU16, Perbill,
 };
 use sp_staking::{
@@ -34,19 +40,23 @@ use sp_staking::{
 use sp_std::prelude::*;
 use std::collections::BTreeMap;
 
+use codec::Decode;
 use frame_election_provider_support::{onchain, ElectionDataProvider, SequentialPhragmen, Weight};
 use pallet_election_provider_multi_phase::{
-	unsigned::MinerConfig, ElectionCompute, QueuedSolution, SolutionAccuracyOf,
+	unsigned::MinerConfig, Call, ElectionCompute, QueuedSolution, SolutionAccuracyOf,
 };
 use pallet_staking::StakerStatus;
+use parking_lot::RwLock;
+use std::sync::Arc;
+
+use frame_support::derive_impl;
 
 use crate::{log, log_current_time};
 
-pub const INIT_TIMESTAMP: u64 = 30_000;
-pub const BLOCK_TIME: u64 = 1000;
-
-type Block = frame_system::mocking::MockBlock<Runtime>;
+pub const INIT_TIMESTAMP: BlockNumber = 30_000;
+pub const BLOCK_TIME: BlockNumber = 1000;
 
+type Block = frame_system::mocking::MockBlockU32<Runtime>;
 type Extrinsic = testing::TestXt<RuntimeCall, ()>;
 
 frame_support::construct_runtime!(
@@ -63,38 +73,26 @@ frame_support::construct_runtime!(
 	}
 );
 
-pub(crate) type AccountId = u128;
-pub(crate) type Nonce = u32;
-pub(crate) type BlockNumber = u64;
+pub(crate) type AccountId = u64;
+pub(crate) type AccountIndex = u32;
+pub(crate) type BlockNumber = u32;
 pub(crate) type Balance = u64;
-pub(crate) type VoterIndex = u32;
+pub(crate) type VoterIndex = u16;
 pub(crate) type TargetIndex = u16;
-pub(crate) type Moment = u64;
+pub(crate) type Moment = u32;
 
+#[derive_impl(frame_system::config_preludes::TestDefaultConfig as frame_system::DefaultConfig)]
 impl frame_system::Config for Runtime {
-	type BaseCallFilter = traits::Everything;
-	type BlockWeights = BlockWeights;
-	type BlockLength = ();
-	type DbWeight = ();
-	type RuntimeOrigin = RuntimeOrigin;
-	type Nonce = Nonce;
-	type RuntimeCall = RuntimeCall;
-	type Hash = H256;
-	type Hashing = sp_runtime::traits::BlakeTwo256;
-	type AccountId = AccountId;
-	type Lookup = IdentityLookup<Self::AccountId>;
 	type Block = Block;
+	type BlockHashCount = ConstU32<10>;
+	type BaseCallFilter = frame_support::traits::Everything;
+	type RuntimeOrigin = RuntimeOrigin;
+	type RuntimeCall = RuntimeCall;
 	type RuntimeEvent = RuntimeEvent;
-	type BlockHashCount = ();
-	type Version = ();
 	type PalletInfo = PalletInfo;
-	type AccountData = pallet_balances::AccountData<Balance>;
-	type OnNewAccount = ();
-	type OnKilledAccount = ();
-	type SystemWeightInfo = ();
-	type SS58Prefix = ();
 	type OnSetCode = ();
-	type MaxConsumers = traits::ConstU32<16>;
+
+	type AccountData = pallet_balances::AccountData<Balance>;
 }
 
 const NORMAL_DISPATCH_RATIO: Perbill = Perbill::from_percent(75);
@@ -126,13 +124,13 @@ impl pallet_balances::Config for Runtime {
 impl pallet_timestamp::Config for Runtime {
 	type Moment = Moment;
 	type OnTimestampSet = ();
-	type MinimumPeriod = traits::ConstU64<5>;
+	type MinimumPeriod = traits::ConstU32<5>;
 	type WeightInfo = ();
 }
 
 parameter_types! {
-	pub static Period: BlockNumber = 30;
-	pub static Offset: BlockNumber = 0;
+	pub static Period: u32 = 30;
+	pub static Offset: u32 = 0;
 }
 
 sp_runtime::impl_opaque_keys! {
@@ -180,6 +178,7 @@ parameter_types! {
 	pub static MinerMaxLength: u32 = 256;
 	pub static MinerMaxWeight: Weight = BlockWeights::get().max_block;
 	pub static TransactionPriority: transaction_validity::TransactionPriority = 1;
+	#[derive(Debug)]
 	pub static MaxWinners: u32 = 100;
 	pub static MaxVotesPerVoter: u32 = 16;
 	pub static MaxNominations: u32 = 16;
@@ -479,17 +478,18 @@ impl Default for ExtBuilder {
 }
 
 impl ExtBuilder {
-	pub fn build(self) -> sp_io::TestExternalities {
+	pub fn build(&self) -> sp_io::TestExternalities {
 		sp_tracing::try_init_simple();
 		let mut storage =
 			frame_system::GenesisConfig::<Runtime>::default().build_storage().unwrap();
 
-		let _ =
-			pallet_balances::GenesisConfig::<Runtime> { balances: self.balances_builder.balances }
-				.assimilate_storage(&mut storage);
+		let _ = pallet_balances::GenesisConfig::<Runtime> {
+			balances: self.balances_builder.balances.clone(),
+		}
+		.assimilate_storage(&mut storage);
 
 		let mut stakers = self.staking_builder.stakers.clone();
-		self.staking_builder.status.into_iter().for_each(|(stash, status)| {
+		self.staking_builder.status.clone().into_iter().for_each(|(stash, status)| {
 			let (_, _, _, ref mut prev_status) = stakers
 				.iter_mut()
 				.find(|s| s.0 == stash)
@@ -497,7 +497,7 @@ impl ExtBuilder {
 			*prev_status = status;
 		});
 		// replaced any of the stakes if needed.
-		self.staking_builder.stakes.into_iter().for_each(|(stash, stake)| {
+		self.staking_builder.stakes.clone().into_iter().for_each(|(stash, stake)| {
 			let (_, _, ref mut prev_stake, _) = stakers
 				.iter_mut()
 				.find(|s| s.0 == stash)
@@ -532,12 +532,13 @@ impl ExtBuilder {
 		ext.execute_with(|| {
 			System::set_block_number(1);
 			Session::on_initialize(1);
-			<Staking as Hooks<u64>>::on_initialize(1);
+			<Staking as Hooks<u32>>::on_initialize(1);
 			Timestamp::set_timestamp(INIT_TIMESTAMP);
 		});
 
 		ext
 	}
+
 	pub fn staking(mut self, builder: StakingExtBuilder) -> Self {
 		self.staking_builder = builder;
 		self
@@ -553,8 +554,33 @@ impl ExtBuilder {
 		self
 	}
 
+	pub fn build_offchainify(
+		self,
+	) -> (sp_io::TestExternalities, Arc<RwLock<PoolState>>, Arc<RwLock<OffchainState>>) {
+		// add offchain and pool externality extensions.
+		let mut ext = self.build();
+		let (offchain, offchain_state) = TestOffchainExt::new();
+		let (pool, pool_state) = TestTransactionPoolExt::new();
+
+		ext.register_extension(OffchainDbExt::new(offchain.clone()));
+		ext.register_extension(OffchainWorkerExt::new(offchain));
+		ext.register_extension(TransactionPoolExt::new(pool));
+
+		(ext, pool_state, offchain_state)
+	}
+
 	pub fn build_and_execute(self, test: impl FnOnce() -> ()) {
-		self.build().execute_with(test)
+		let mut ext = self.build();
+		ext.execute_with(test);
+
+		#[cfg(feature = "try-runtime")]
+		ext.execute_with(|| {
+			let bn = System::block_number();
+
+			assert_ok!(<MultiPhase as Hooks<u64>>::try_state(bn));
+			assert_ok!(<Staking as Hooks<u64>>::try_state(bn));
+			assert_ok!(<Session as Hooks<u64>>::try_state(bn));
+		});
 	}
 }
 
@@ -586,18 +612,67 @@ pub fn roll_to(n: BlockNumber, delay_solution: bool) {
 	}
 }
 
+// Progress to given block, triggering session and era changes as we progress and ensuring that
+// there is a solution queued when expected.
+pub fn roll_to_with_ocw(n: BlockNumber, pool: Arc<RwLock<PoolState>>, delay_solution: bool) {
+	for b in (System::block_number()) + 1..=n {
+		System::set_block_number(b);
+		Session::on_initialize(b);
+		Timestamp::set_timestamp(System::block_number() * BLOCK_TIME + INIT_TIMESTAMP);
+
+		ElectionProviderMultiPhase::on_initialize(b);
+		ElectionProviderMultiPhase::offchain_worker(b);
+
+		if !delay_solution && pool.read().transactions.len() > 0 {
+			// decode submit_unsigned callable that may be queued in the pool by ocw. skip all
+			// other extrinsics in the pool.
+			for encoded in &pool.read().transactions {
+				let extrinsic = Extrinsic::decode(&mut &encoded[..]).unwrap();
+
+				let _ = match extrinsic.call {
+					RuntimeCall::ElectionProviderMultiPhase(
+						call @ Call::submit_unsigned { .. },
+					) => {
+						// call submit_unsigned callable in OCW pool.
+						crate::assert_ok!(call.dispatch_bypass_filter(RuntimeOrigin::none()));
+					},
+					_ => (),
+				};
+			}
+
+			pool.try_write().unwrap().transactions.clear();
+		}
+
+		Staking::on_initialize(b);
+		if b != n {
+			Staking::on_finalize(System::block_number());
+		}
+
+		log_current_time();
+	}
+}
+// helper to progress one block ahead.
+pub fn roll_one(pool: Arc<RwLock<PoolState>>, delay_solution: bool) {
+	let bn = System::block_number().saturating_add(1);
+	roll_to_with_ocw(bn, pool, delay_solution);
+}
+
 /// Progresses from the current block number (whatever that may be) to the block where the session
 /// `session_index` starts.
-pub(crate) fn start_session(session_index: SessionIndex, delay_solution: bool) {
-	let end: u64 = if Offset::get().is_zero() {
-		Period::get() * (session_index as u64)
+pub(crate) fn start_session(
+	session_index: SessionIndex,
+	pool: Arc<RwLock<PoolState>>,
+	delay_solution: bool,
+) {
+	let end = if Offset::get().is_zero() {
+		Period::get() * session_index
 	} else {
-		Offset::get() * (session_index as u64) + Period::get() * (session_index as u64)
+		Offset::get() * session_index + Period::get() * session_index
 	};
 
 	assert!(end >= System::block_number());
 
-	roll_to(end, delay_solution);
+	roll_to_with_ocw(end, pool, delay_solution);
 
 	// session must have progressed properly.
 	assert_eq!(
@@ -610,29 +685,35 @@ pub(crate) fn start_session(session_index: SessionIndex, delay_solution: bool) {
 }
 
 /// Go one session forward.
-pub(crate) fn advance_session() {
+pub(crate) fn advance_session(pool: Arc<RwLock<PoolState>>) {
 	let current_index = Session::current_index();
-	start_session(current_index + 1, false);
+	start_session(current_index + 1, pool, false);
 }
 
-pub(crate) fn advance_session_delayed_solution() {
+pub(crate) fn advance_session_delayed_solution(pool: Arc<RwLock<PoolState>>) {
 	let current_index = Session::current_index();
-	start_session(current_index + 1, true);
+	start_session(current_index + 1, pool, true);
 }
 
-pub(crate) fn start_next_active_era() -> Result<(), ()> {
-	start_active_era(active_era() + 1, false)
+pub(crate) fn start_next_active_era(pool: Arc<RwLock<PoolState>>) -> Result<(), ()> {
+	start_active_era(active_era() + 1, pool, false)
 }
 
-pub(crate) fn start_next_active_era_delayed_solution() -> Result<(), ()> {
-	start_active_era(active_era() + 1, true)
+pub(crate) fn start_next_active_era_delayed_solution(
+	pool: Arc<RwLock<PoolState>>,
+) -> Result<(), ()> {
+	start_active_era(active_era() + 1, pool, true)
 }
 
 /// Progress until the given era.
-pub(crate) fn start_active_era(era_index: EraIndex, delay_solution: bool) -> Result<(), ()> {
+pub(crate) fn start_active_era(
+	era_index: EraIndex,
+	pool: Arc<RwLock<PoolState>>,
+	delay_solution: bool,
+) -> Result<(), ()> {
 	let era_before = current_era();
 
-	start_session((era_index * <SessionsPerEra as Get<u32>>::get()).into(), delay_solution);
+	start_session((era_index * <SessionsPerEra as Get<u32>>::get()).into(), pool, delay_solution);
 
 	log!(
 		info,
@@ -745,7 +826,7 @@ pub(crate) fn slash_through_offending_threshold() {
 
 // Slashes a percentage of the active nominators that haven't been slashed yet, with
 // a minimum of 1 validator slash.
-pub(crate) fn slash_percentage(percentage: Perbill) -> Vec<u128> {
+pub(crate) fn slash_percentage(percentage: Perbill) -> Vec<AccountId> {
 	let validators = Session::validators();
 	let mut remaining_slashes = (percentage * validators.len() as u32).max(1);
 	let mut slashed = vec![];
@@ -781,3 +862,17 @@ pub(crate) fn staking_events() -> Vec<pallet_staking::Event<Runtime>> {
 		.filter_map(|e| if let RuntimeEvent::Staking(inner) = e { Some(inner) } else { None })
 		.collect::<Vec<_>>()
 }
+
+pub(crate) fn epm_events() -> Vec<pallet_election_provider_multi_phase::Event<Runtime>> {
+	System::events()
+		.into_iter()
+		.map(|r| r.event)
+		.filter_map(|e| {
+			if let RuntimeEvent::ElectionProviderMultiPhase(inner) = e {
+				Some(inner)
+			} else {
+				None
+			}
+		})
+		.collect::<Vec<_>>()
+}