Repository: paritytech/substrate
File: frame/election-provider-multi-phase/test-staking-e2e/src/mock.rs (Rust)

Code before:
		self.balances_builder = builder;
		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() -> ()) {
		let mut ext = self.build();
		ext.execute_with(test);

		#[cfg(feature = "try-runtime")]
		ext.execute_with(|| {
			let bm = 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));
		});
	}
}

// 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(n: BlockNumber, 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);

		// TODO(gpestana): implement a realistic OCW worker insted of simulating it
		// https://github.com/paritytech/substrate/issues/13589
		// if there's no solution queued and the solution should not be delayed, try mining and
		// queue a solution.
		if ElectionProviderMultiPhase::current_phase().is_signed() && !delay_solution {
			let _ = try_queue_solution(ElectionCompute::Signed).map_err(|e| {
				log!(info, "failed to mine/queue solution: {:?}", e);
			});
		}

Review comment (suggestion):
```suggestion
			let bn = System::block_number();
```

Code after:
		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() -> ()) {
		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));
		});
	}
}

// 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(n: BlockNumber, 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);

		// TODO(gpestana): implement a realistic OCW worker insted of simulating it
		// https://github.com/paritytech/substrate/issues/13589
		// if there's no solution queued and the solution should not be delayed, try mining and
		// queue a solution.
		if ElectionProviderMultiPhase::current_phase().is_signed() && !delay_solution {
			let _ = try_queue_solution(ElectionCompute::Signed).map_err(|e| {
				log!(info, "failed to mine/queue solution: {:?}", e);
			});
		}
		ElectionProviderMultiPhase::on_initialize(b);

		Staking::on_initialize(b);
		if b != n {
			Staking::on_finalize(System::block_number());