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

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

		log_current_time();
	}
}

// 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 {
			// build and submit extrinsics queued in the pool by ocw.
			let encoded = pool.read().transactions[0].clone();
			let extrinsic = Extrinsic::decode(&mut &*encoded).unwrap();
			let call = match extrinsic.call {
				RuntimeCall::ElectionProviderMultiPhase(call @ Call::submit_unsigned { .. }) =>
					call,
				_ => panic!(
					"unexpected error when decoding the submit_unsigned extrinsic from the 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

Review comment (suggestion):
maybe something like this would be simpler? 
```
assert!(matches!(extrinsic.call), RuntimeCall::ElectionProviderMultiPhase(Call::submit_unsigned { .. }))
```

Code after:
		log_current_time();
	}
}

// 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.