diff --git a/substrate/frame/support/src/storage/storage_noop_guard.rs b/substrate/frame/support/src/storage/storage_noop_guard.rs index d00e6e18ecc4..c4d40fa99a35 100644 --- a/substrate/frame/support/src/storage/storage_noop_guard.rs +++ b/substrate/frame/support/src/storage/storage_noop_guard.rs @@ -37,15 +37,38 @@ /// }); /// ``` #[must_use] -pub struct StorageNoopGuard(sp_std::vec::Vec); +pub struct StorageNoopGuard<'a> { + storage_root: sp_std::vec::Vec, + 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() { @@ -53,8 +76,9 @@ impl Drop for StorageNoopGuard { } 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"); + }); + } }