Improve post and pre migration checks for pallet-membership

frame/membership/src/lib.rs

@@ -790,9 +790,16 @@ mod tests {
 	fn migration_v4() {
 		new_test_ext().execute_with(|| {
 			use frame_support::traits::PalletInfo;
-			let old_pallet_name =
+			let old_pallet_name = "OldMembership";
+			let new_pallet_name =
 				<Test as frame_system::Config>::PalletInfo::name::<Membership>().unwrap();
-			let new_pallet_name = "NewMembership";
+
+			frame_support::storage::migration::move_pallet(
+				new_pallet_name.as_bytes(),
+				old_pallet_name.as_bytes(),
+			);
+
+			StorageVersion::new(0).put::<Membership>();
 
 			crate::migrations::v4::pre_migrate::<Membership, _>(old_pallet_name, new_pallet_name);
 			crate::migrations::v4::migrate::<Test, Membership, _>(old_pallet_name, new_pallet_name);

frame/membership/src/migrations/v4.rs

@@ -15,8 +15,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use sp_core::hexdisplay::HexDisplay;
-use sp_io::{hashing::twox_128, storage};
+use sp_io::hashing::twox_128;
 
 use frame_support::{
 	traits::{
@@ -85,28 +84,22 @@ pub fn pre_migrate<P: GetStorageVersion, N: AsRef<str>>(old_pallet_name: N, new_
 	let new_pallet_name = new_pallet_name.as_ref();
 	log_migration("pre-migration", old_pallet_name, new_pallet_name);
 
-	let old_pallet_prefix = twox_128(old_pallet_name.as_bytes());
-	assert!(storage::next_key(&old_pallet_prefix)
-		.map_or(true, |next_key| next_key.starts_with(&old_pallet_prefix)));
+	if new_pallet_name == old_pallet_name {
+		return
+	}
 
 	let new_pallet_prefix = twox_128(new_pallet_name.as_bytes());
-	let storage_version_key =
-		[&new_pallet_prefix, &twox_128(STORAGE_VERSION_STORAGE_KEY_POSTFIX)[..]].concat();
-	// ensure nothing is stored in the new prefix.
-	assert!(
-		storage::next_key(&new_pallet_prefix).map_or(
-			// either nothing is there
-			true,
-			// or we ensure that it has no common prefix with twox_128(new),
-			// or isn't the storage version that is already stored using the pallet name
-			|next_key| {
-				!next_key.starts_with(&new_pallet_prefix) || next_key == storage_version_key
-			},
-		),
-		"unexpected next_key({}) = {:?}",
-		new_pallet_name,
-		HexDisplay::from(&storage::next_key(&new_pallet_prefix).unwrap()),
+	let storage_version_key = twox_128(STORAGE_VERSION_STORAGE_KEY_POSTFIX);
+
+	let mut new_pallet_prefix_iter = frame_support::storage::KeyPrefixIterator::new(
+		new_pallet_prefix.to_vec(),
+		new_pallet_prefix.to_vec(),
+		|key| Ok(key.to_vec()),
 	);
+
+	// Ensure nothing except maybe the storage_version_key is stored in the new prefix.
+	assert!(new_pallet_prefix_iter.all(|key| key == storage_version_key));
+
 	assert!(<P as GetStorageVersion>::on_chain_storage_version() < 4);
 }
 
@@ -119,26 +112,27 @@ pub fn post_migrate<P: GetStorageVersion, N: AsRef<str>>(old_pallet_name: N, new
 	let new_pallet_name = new_pallet_name.as_ref();
 	log_migration("post-migration", old_pallet_name, new_pallet_name);
 
-	let old_pallet_prefix = twox_128(old_pallet_name.as_bytes());
-	#[cfg(test)]
-	{
-		let storage_version_key =
-			[&old_pallet_prefix, &twox_128(STORAGE_VERSION_STORAGE_KEY_POSTFIX)[..]].concat();
-		assert!(storage::next_key(&old_pallet_prefix)
-			.map_or(true, |next_key| !next_key.starts_with(&old_pallet_prefix) ||
-				next_key == storage_version_key));
-	}
-	#[cfg(not(test))]
-	{
-		// Assert that nothing remains at the old prefix
-		assert!(storage::next_key(&old_pallet_prefix)
-			.map_or(true, |next_key| !next_key.starts_with(&old_pallet_prefix)));
+	if new_pallet_name == old_pallet_name {
+		return
 	}
 
+	// Assert that nothing remains at the old prefix.
+	let old_pallet_prefix = twox_128(old_pallet_name.as_bytes());
+	let old_pallet_prefix_iter = frame_support::storage::KeyPrefixIterator::new(
+		old_pallet_prefix.to_vec(),
+		old_pallet_prefix.to_vec(),
+		|_| Ok(()),
+	);
+	assert_eq!(old_pallet_prefix_iter.count(), 0);
+
+	// NOTE: storage_version_key is already in the new prefix.
 	let new_pallet_prefix = twox_128(new_pallet_name.as_bytes());
-	// Assert that the storages have been moved to the new prefix
-	assert!(storage::next_key(&new_pallet_prefix)
-		.map_or(true, |next_key| next_key.starts_with(&new_pallet_prefix)));
+	let new_pallet_prefix_iter = frame_support::storage::KeyPrefixIterator::new(
+		new_pallet_prefix.to_vec(),
+		new_pallet_prefix.to_vec(),
+		|_| Ok(()),
+	);
+	assert!(new_pallet_prefix_iter.count() >= 1);
 
 	assert_eq!(<P as GetStorageVersion>::on_chain_storage_version(), 4);
 }