Add custom error message for StorageNoopGuard

substrate/frame/support/src/storage/storage_noop_guard.rs

@@ -37,24 +37,48 @@
 /// });
 /// ```
 #[must_use]
-pub struct StorageNoopGuard(sp_std::vec::Vec<u8>);
+pub struct StorageNoopGuard<'a> {
+	storage_root: sp_std::vec::Vec<u8>,
+	error_message: &'a str,
+}
 
-impl Default for StorageNoopGuard {
+impl<'a> Default for StorageNoopGuard<'a> {
 	fn default() -> Self {
-		Self(sp_io::storage::root(sp_runtime::StateVersion::V1))
+		Self {
+			storage_root: sp_io::storage::root(sp_runtime::StateVersion::V1),
+			error_message: "`StorageNoopGuard` detected an attempted storage change.",
+		}
+	}
+}
+
+impl<'a> StorageNoopGuard<'a> {
+	/// Alias to `default()`.
+	pub fn new() -> Self {
+		Self::default()
+	}
+
+	/// Creates a new [`StorageNoopGuard`] with a custom error message.
+	pub fn from_error_message(error_message: &'a str) -> Self {
+		Self { storage_root: sp_io::storage::root(sp_runtime::StateVersion::V1), error_message }
+	}
+
+	/// Sets a custom error message for a [`StorageNoopGuard`].
+	pub fn set_error_message(&mut self, error_message: &'a str) {
+		self.error_message = error_message;
 	}
 }
 
-impl Drop for StorageNoopGuard {
+impl<'a> Drop for StorageNoopGuard<'a> {
 	fn drop(&mut self) {
 		// No need to double panic, eg. inside a test assertion failure.
 		if sp_std::thread::panicking() {
 			return
 		}
 		assert_eq!(
 			sp_io::storage::root(sp_runtime::StateVersion::V1),
-			self.0,
-			"StorageNoopGuard detected wrongful storage changes.",
+			self.storage_root,
+			"{}",
+			self.error_message,
 		);
 	}
 }
@@ -65,7 +89,7 @@ mod tests {
 	use sp_io::TestExternalities;
 
 	#[test]
-	#[should_panic(expected = "StorageNoopGuard detected wrongful storage changes.")]
+	#[should_panic(expected = "`StorageNoopGuard` detected an attempted storage change.")]
 	fn storage_noop_guard_panics_on_changed() {
 		TestExternalities::default().execute_with(|| {
 			let _guard = StorageNoopGuard::default();
@@ -83,7 +107,7 @@ mod tests {
 	}
 
 	#[test]
-	#[should_panic(expected = "StorageNoopGuard detected wrongful storage changes.")]
+	#[should_panic(expected = "`StorageNoopGuard` detected an attempted storage change.")]
 	fn storage_noop_guard_panics_on_early_drop() {
 		TestExternalities::default().execute_with(|| {
 			let guard = StorageNoopGuard::default();
@@ -111,4 +135,34 @@ mod tests {
 			panic!("Something else");
 		});
 	}
+
+	#[test]
+	#[should_panic(expected = "`StorageNoopGuard` found unexpected storage changes.")]
+	fn storage_noop_guard_panics_created_from_error_message() {
+		TestExternalities::default().execute_with(|| {
+			let _guard = StorageNoopGuard::from_error_message(
+				"`StorageNoopGuard` found unexpected storage changes.",
+			);
+			frame_support::storage::unhashed::put(b"key", b"value");
+		});
+	}
+
+	#[test]
+	#[should_panic(expected = "`StorageNoopGuard` found unexpected storage changes.")]
+	fn storage_noop_guard_panics_with_set_error_message() {
+		TestExternalities::default().execute_with(|| {
+			let mut guard = StorageNoopGuard::default();
+			guard.set_error_message("`StorageNoopGuard` found unexpected storage changes.");
+			frame_support::storage::unhashed::put(b"key", b"value");
+		});
+	}
+
+	#[test]
+	#[should_panic(expected = "`StorageNoopGuard` detected an attempted storage change.")]
+	fn storage_noop_guard_panics_new_alias() {
+		TestExternalities::default().execute_with(|| {
+			let _guard = StorageNoopGuard::new();
+			frame_support::storage::unhashed::put(b"key", b"value");
+		});
+	}
 }