From 77eba1351a1a2f74360b61143affc27983f80e4e Mon Sep 17 00:00:00 2001
From: Parth <desaiparth08@gmail.com>
Date: Fri, 21 May 2021 15:51:12 +0530
Subject: [PATCH] add ExternalitiesExt for TestExternalities (#8877)

---
 .../primitives/state-machine/src/testing.rs   | 22 ++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/substrate/primitives/state-machine/src/testing.rs b/substrate/primitives/state-machine/src/testing.rs
index f4b0cb6592..250c2fd4e9 100644
--- a/substrate/primitives/state-machine/src/testing.rs
+++ b/substrate/primitives/state-machine/src/testing.rs
@@ -41,7 +41,7 @@ use sp_core::{
 	traits::TaskExecutorExt,
 	testing::TaskExecutor,
 };
-use sp_externalities::{Extensions, Extension};
+use sp_externalities::{Extensions, Extension, ExtensionStore};
 
 /// Simple HashMap-based Externalities impl.
 pub struct TestExternalities<H: Hasher, N: ChangesTrieBlockNumber = u64>
@@ -274,6 +274,26 @@ impl<H, N> sp_externalities::ExtensionStore for TestExternalities<H, N> where
 	}
 }
 
+impl<H, N> sp_externalities::ExternalitiesExt for TestExternalities<H, N>
+	where
+		H: Hasher,
+		H::Out: Ord + codec::Codec,
+		N: ChangesTrieBlockNumber,
+{
+	fn extension<T: Any + Extension>(&mut self) -> Option<&mut T> {
+		self.extension_by_type_id(TypeId::of::<T>())
+			.and_then(<dyn Any>::downcast_mut)
+	}
+
+	fn register_extension<T: Extension>(&mut self, ext: T) -> Result<(), sp_externalities::Error> {
+		self.register_extension_with_type_id(TypeId::of::<T>(), Box::new(ext))
+	}
+
+	fn deregister_extension<T: Extension>(&mut self) -> Result<(), sp_externalities::Error> {
+		self.deregister_extension_by_type_id(TypeId::of::<T>())
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	use super::*;