Repository: paritytech/substrate
File: frame/elections-phragmen/src/migrations/unlock_and_unreserve_all_funds.rs (Rust)

Code before:
	/// 2. Unreserves the deposited funds for each account.
	/// 3. Unlocks the staked funds for each account.
	fn on_runtime_upgrade() -> frame_support::weights::Weight {
		// Get staked and deposited balances as reported by this pallet.
		let (account_deposited_sums, account_staked_sums, initial_reads_weight) =
			Self::get_account_deposited_and_staked_sums();

		// Deposited funds need to be unreserved.
		for (account, unreserve_amount) in account_deposited_sums {
			if unreserve_amount.is_zero() {
				continue
			}
			T::Currency::unreserve(&account, unreserve_amount);
		}

		// Staked funds need to be unlocked.
		for (account, amount) in account_staked_sums {
			if amount.is_zero() {
				continue
			}
			// Question for reviewers:
			// remove_lock triggers a few
			// `ERROR main runtime::system: Logic error: Unexpected underflow in reducing consumer`
			// errors. are they safe to ignore?
			//
			// I have verified that a non-zero lock always exists before remove_lock is called.
			T::Currency::remove_lock(T::PalletId::get(), &account);
		}

		// Question for reviewers: how do I know the weight of the unreserve & remove_lock calls?
		RocksDbWeight::get().reads_writes(1, 1) + initial_reads_weight
	}

	/// Performs post-upgrade sanity checks:
	///
	/// 1. All expected locks were removed after the migration.
	/// 2. There are no 'hanging' locks for this pallet in Balances.
	/// 3. The reserved balance for each account has been reduced by the expected amount.
	#[cfg(feature = "try-runtime")]
	fn post_upgrade(pre_migration_data_bytes: Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
		use sp_runtime::Saturating;

		let pre_migration_data = PreMigrationData::<T>::decode(&mut &pre_migration_data_bytes[..])
			.map_err(|_| "Failed to decode pre-migration data")?;

		// Get staked and deposited balances as reported by this pallet.
		let (account_deposited_sums, account_staked_sums, _) =
			Self::get_account_deposited_and_staked_sums();

		// Check that all expected locks were removed.
		account_staked_sums.iter().for_each(|(acc, _)| {

Review comment (refactor):
How about we also test the pre/post migration checks in the tests? It may require some refactoring though, something along the lines of

```rust
impl<T: crate::Config + pallet_balances::Config> UnlockAndUnreserveAllFunds<T> {
 fn do_pre_upgrade() -> Result<Vec<u8>, sp_runtime::TryRuntimeError> {
  // impl here
 }

 fn do_post_upgrade(pre_migration_data_bytes: Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
  // impl here
 }
}

impl<T: crate::Config + pallet_balances::Config> OnRuntimeUpgrade for UnlockAndUnreserveAllFunds<T> { 
 #[cfg(feature = "try-runtime")]
 fn pre_upgrade() -> Result<Vec<u8>, sp_runtime::TryRuntimeError> {
   Self::do_pre_upgrade()
 } 

 #[cfg(feature = "try-runtime")]
 fn post_upgrade(pre_migration_data_bytes: Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
   Self::do_post_upgrade()
 } 
}
```

And then we can add the checks in the tests:

```rust
#[test]
fn test() {
 // ... 
 let pre_migration_state = UnlockAndUnreserveAllFunds::<Test>::do_pre_upgrade().unwrap();
 UnlockAndUnreserveAllFunds::<Test>::on_runtime_upgrade();
 assert_ok!(UnlockAndUnreserveAllFunds::<Test>::do_post_upgrade(pre_migration_state));
 // ... 
}
```

Code after:
		for (account, actual_reserved_before) in account_reserved_before {
			let actual_reserved_after = T::Currency::reserved_balance(&account);
			let expected_amount_deducted = *account_deposited_sums
				.get(&account)
				.unwrap_or(&Zero::zero())
				// .min here to handle bugged deposits where actual_reserved_before is less than the
				// amount the pallet reports is reserved
				.min(&actual_reserved_before);
			let expected_reserved_after =
				actual_reserved_before.saturating_sub(expected_amount_deducted);
			assert!(
				actual_reserved_after == expected_reserved_after,
				"Reserved balance for {:?} is incorrect. actual before: {:?}, actual after, {:?}, expected deducted: {:?}",
				account,
				actual_reserved_before,
				actual_reserved_after,
				expected_amount_deducted,
			);
		}

		Ok(())
	}
}

#[cfg(all(feature = "try-runtime", test))]
mod test {
	use super::*;
	use crate::{
		tests::{ExtBuilder, Test},
		Candidates, Members, RunnersUp, SeatHolder, Voter, Voting,
	};
	use frame_support::{
		assert_ok,
		traits::{Currency, OnRuntimeUpgrade, ReservableCurrency, WithdrawReasons},
	};

	#[test]
	fn unreserve_works_for_candidate() {
		let candidate = 10;
		let deposit = 100;
		let initial_reserved = 15;
		let initial_balance = 100_000;
		ExtBuilder::default().build_and_execute(|| {
			// Set up initial state.
			<Test as crate::Config>::Currency::make_free_balance_be(&candidate, initial_balance);
			assert_ok!(<Test as crate::Config>::Currency::reserve(&candidate, initial_reserved));
			Candidates::<Test>::set(vec![(candidate, deposit)]);
			assert_ok!(<Test as crate::Config>::Currency::reserve(&candidate, deposit));

			// Sanity check: ensure initial Balance state was set up correctly.
			assert_eq!(