From 8e90d6a59d4f4917a923ba4de2ecdfe3ecae647d Mon Sep 17 00:00:00 2001
From: wangjj9219 <183318287@qq.com>
Date: Tue, 14 Mar 2023 01:28:23 +0800
Subject: [PATCH] fix migration and tests

---
 asset-registry/src/migrations.rs |  5 +-
 asset-registry/src/tests.rs      | 88 ++++++++++++++++++++++++++++----
 unknown-tokens/src/migrations.rs |  4 +-
 3 files changed, 83 insertions(+), 14 deletions(-)

diff --git a/asset-registry/src/migrations.rs b/asset-registry/src/migrations.rs
index 1ce62e72b..75320e289 100644
--- a/asset-registry/src/migrations.rs
+++ b/asset-registry/src/migrations.rs
@@ -27,8 +27,9 @@ mod v2 {
 
 		weight.saturating_accrue(T::DbWeight::get().reads(1));
 		let old_data =
-			storage_key_iter::<xcm::v2::MultiLocation, T::AssetId, Blake2_128Concat>(module_prefix, storage_prefix)
-				.drain();
+			storage_key_iter::<xcm::v2::MultiLocation, T::AssetId, Twox64Concat>(module_prefix, storage_prefix)
+				.drain()
+				.collect::<sp_std::vec::Vec<_>>();
 
 		for (old_key, value) in old_data {
 			weight.saturating_accrue(T::DbWeight::get().writes(1));
diff --git a/asset-registry/src/tests.rs b/asset-registry/src/tests.rs
index 4a0f53169..d43e34dc1 100644
--- a/asset-registry/src/tests.rs
+++ b/asset-registry/src/tests.rs
@@ -601,28 +601,71 @@ fn from_unversioned_to_v2_storage() {
 		// StorageVersion is 0 before migration
 		assert_eq!(StorageVersion::get::<Pallet<para::Runtime>>(), 0);
 
-		// V2 storage key
-		let old_key = xcm::v2::MultiLocation::new(
+		// V2 storage
+		let old_multilocation_0 = xcm::v2::MultiLocation::new(
 			0,
 			xcm::v2::Junctions::X1(xcm::v2::Junction::GeneralKey(vec![0].try_into().unwrap())),
-		)
-		.encode();
+		);
+		let old_multilocation_1 = xcm::v2::MultiLocation::new(
+			1,
+			xcm::v2::Junctions::X2(
+				xcm::v2::Junction::Parachain(2096),
+				xcm::v2::Junction::GeneralKey(vec![0, 0, 0, 0, 0, 0, 0, 0, 0].try_into().unwrap()),
+			),
+		);
+		let old_multilocation_2 = xcm::v2::MultiLocation::new(
+			1,
+			xcm::v2::Junctions::X2(
+				xcm::v2::Junction::Parachain(2096),
+				xcm::v2::Junction::GeneralKey(vec![1, 1].try_into().unwrap()),
+			),
+		);
 
-		let asset_id: para::ParaAssetId = 5u32;
+		let asset_id_0: para::ParaAssetId = 5u32;
+		let asset_id_1: para::ParaAssetId = 6u32;
+		let asset_id_2: para::ParaAssetId = 7u32;
 
 		// Store raw xcm::v2 data
 		put_storage_value(
 			module_prefix,
 			storage_prefix,
-			&Blake2_128Concat::hash(&old_key),
-			asset_id,
+			&Twox64Concat::hash(&old_multilocation_0.encode()),
+			asset_id_0,
+		);
+		put_storage_value(
+			module_prefix,
+			storage_prefix,
+			&Twox64Concat::hash(&old_multilocation_1.encode()),
+			asset_id_1,
+		);
+		put_storage_value(
+			module_prefix,
+			storage_prefix,
+			&Twox64Concat::hash(&old_multilocation_2.encode()),
+			asset_id_2,
 		);
 
 		// V3 storage key
-		let new_key = MultiLocation::new(0, X1(Junction::from(BoundedVec::try_from(vec![0]).unwrap())));
+		let new_multilocation_0 = MultiLocation::new(0, X1(Junction::from(BoundedVec::try_from(vec![0]).unwrap())));
+		let new_multilocation_1 = MultiLocation::new(
+			1,
+			X2(
+				Parachain(2096),
+				Junction::from(BoundedVec::try_from(vec![0, 0, 0, 0, 0, 0, 0, 0, 0]).unwrap()),
+			),
+		);
+		let new_multilocation_2 = MultiLocation::new(
+			1,
+			X2(
+				Parachain(2096),
+				Junction::from(BoundedVec::try_from(vec![1, 1]).unwrap()),
+			),
+		);
 
 		// Assert new StorageKey still does not exist
-		assert_eq!(AssetRegistry::location_to_asset_id(new_key), None);
+		assert_eq!(AssetRegistry::location_to_asset_id(new_multilocation_0), None);
+		assert_eq!(AssetRegistry::location_to_asset_id(new_multilocation_1), None);
+		assert_eq!(AssetRegistry::location_to_asset_id(new_multilocation_2), None);
 
 		// Run StorageKey migration
 		crate::Migration::<para::Runtime>::on_runtime_upgrade();
@@ -631,13 +674,36 @@ fn from_unversioned_to_v2_storage() {
 		assert_eq!(StorageVersion::get::<Pallet<para::Runtime>>(), 2);
 
 		// Assert the StorageKey exists and has been migrated to xcm::v3
-		assert_eq!(AssetRegistry::location_to_asset_id(new_key), Some(asset_id));
+		assert_eq!(
+			AssetRegistry::location_to_asset_id(new_multilocation_0),
+			Some(asset_id_0)
+		);
+		assert_eq!(
+			AssetRegistry::location_to_asset_id(new_multilocation_1),
+			Some(asset_id_1)
+		);
+		assert_eq!(
+			AssetRegistry::location_to_asset_id(new_multilocation_2),
+			Some(asset_id_2)
+		);
 
 		// Assert the old key does not exist anymore
 		assert!(get_storage_value::<para::ParaAssetId>(
 			module_prefix,
 			storage_prefix,
-			&Blake2_128Concat::hash(&old_key),
+			&Twox64Concat::hash(&old_multilocation_0.encode()),
+		)
+		.is_none());
+		assert!(get_storage_value::<para::ParaAssetId>(
+			module_prefix,
+			storage_prefix,
+			&Twox64Concat::hash(&old_multilocation_1.encode()),
+		)
+		.is_none());
+		assert!(get_storage_value::<para::ParaAssetId>(
+			module_prefix,
+			storage_prefix,
+			&Twox64Concat::hash(&old_multilocation_2.encode()),
 		)
 		.is_none());
 
diff --git a/unknown-tokens/src/migrations.rs b/unknown-tokens/src/migrations.rs
index 84a36584e..69a4fa309 100644
--- a/unknown-tokens/src/migrations.rs
+++ b/unknown-tokens/src/migrations.rs
@@ -57,7 +57,9 @@ mod v2 {
 
 		weight.saturating_accrue(T::DbWeight::get().reads(1));
 
-		let old_data = storage_iter::<u128>(module_prefix, storage_prefix).drain();
+		let old_data = storage_iter::<u128>(module_prefix, storage_prefix)
+			.drain()
+			.collect::<sp_std::vec::Vec<_>>();
 
 		for (raw_k, value) in old_data {
 			let mut full_key = Vec::new();