1use std::{any::{Any, TypeId}, panic::{AssertUnwindSafe, UnwindSafe}};
21
22use crate::{
23 backend::Backend, OverlayedChanges, StorageTransactionCache, ext::Ext, InMemoryBackend,
24 StorageKey, StorageValue,
25 changes_trie::{
26 Configuration as ChangesTrieConfiguration,
27 InMemoryStorage as ChangesTrieInMemoryStorage,
28 BlockNumber as ChangesTrieBlockNumber,
29 State as ChangesTrieState,
30 },
31};
32
33use codec::{Decode, Encode};
34use tetsy_hash_db::Hasher;
35use tet_core::{
36 offchain::testing::TestPersistentOffchainDB,
37 storage::{
38 well_known_keys::{CHANGES_TRIE_CONFIG, CODE, HEAP_PAGES, is_child_storage_key},
39 Storage,
40 },
41 traits::TaskExecutorExt,
42 testing::TaskExecutor,
43};
44use externalities::{Extensions, Extension};
45
46pub struct TestExternalities<H: Hasher, N: ChangesTrieBlockNumber = u64>
48where
49 H::Out: codec::Codec + Ord,
50{
51 overlay: OverlayedChanges,
52 offchain_db: TestPersistentOffchainDB,
53 storage_transaction_cache: StorageTransactionCache<
54 <InMemoryBackend<H> as Backend<H>>::Transaction, H, N
55 >,
56 backend: InMemoryBackend<H>,
57 changes_trie_config: Option<ChangesTrieConfiguration>,
58 changes_trie_storage: ChangesTrieInMemoryStorage<H, N>,
59 extensions: Extensions,
60}
61
62impl<H: Hasher, N: ChangesTrieBlockNumber> TestExternalities<H, N>
63 where
64 H::Out: Ord + 'static + codec::Codec
65{
66 pub fn ext(&mut self) -> Ext<H, N, InMemoryBackend<H>> {
68 Ext::new(
69 &mut self.overlay,
70 &mut self.storage_transaction_cache,
71 &self.backend,
72 match self.changes_trie_config.clone() {
73 Some(config) => Some(ChangesTrieState {
74 config,
75 zero: 0.into(),
76 storage: &self.changes_trie_storage,
77 }),
78 None => None,
79 },
80 Some(&mut self.extensions),
81 )
82 }
83
84 pub fn new(storage: Storage) -> Self {
86 Self::new_with_code(&[], storage)
87 }
88
89 pub fn new_empty() -> Self {
91 Self::new_with_code(&[], Storage::default())
92 }
93
94 pub fn new_with_code(code: &[u8], mut storage: Storage) -> Self {
96 let mut overlay = OverlayedChanges::default();
97 let changes_trie_config = storage.top.get(CHANGES_TRIE_CONFIG)
98 .and_then(|v| Decode::decode(&mut &v[..]).ok());
99 overlay.set_collect_extrinsics(changes_trie_config.is_some());
100
101 assert!(storage.top.keys().all(|key| !is_child_storage_key(key)));
102 assert!(storage.children_default.keys().all(|key| is_child_storage_key(key)));
103
104 storage.top.insert(HEAP_PAGES.to_vec(), 8u64.encode());
105 storage.top.insert(CODE.to_vec(), code.to_vec());
106
107 let mut extensions = Extensions::default();
108 extensions.register(TaskExecutorExt::new(TaskExecutor::new()));
109
110 let offchain_db = TestPersistentOffchainDB::new();
111
112 TestExternalities {
113 overlay,
114 offchain_db,
115 changes_trie_config,
116 extensions,
117 changes_trie_storage: ChangesTrieInMemoryStorage::new(),
118 backend: storage.into(),
119 storage_transaction_cache: Default::default(),
120 }
121 }
122
123 pub fn overlayed_changes(&self) -> &OverlayedChanges {
125 &self.overlay
126 }
127
128 pub fn persist_offchain_overlay(&mut self) {
130 self.offchain_db.apply_offchain_changes(self.overlay.offchain_drain_committed());
131 }
132
133 pub fn offchain_db(&self) -> TestPersistentOffchainDB {
135 self.offchain_db.clone()
136 }
137
138 pub fn insert(&mut self, k: StorageKey, v: StorageValue) {
140 self.backend.insert(vec![(None, vec![(k, Some(v))])]);
141 }
142
143 pub fn register_extension<E: Any + Extension>(&mut self, ext: E) {
145 self.extensions.register(ext);
146 }
147
148 pub fn changes_trie_storage(&mut self) -> &mut ChangesTrieInMemoryStorage<H, N> {
150 &mut self.changes_trie_storage
151 }
152
153 fn as_backend(&self) -> InMemoryBackend<H> {
158 let top: Vec<_> = self.overlay.changes()
159 .map(|(k, v)| (k.clone(), v.value().cloned()))
160 .collect();
161 let mut transaction = vec![(None, top)];
162
163 for (child_changes, child_info) in self.overlay.children() {
164 transaction.push((
165 Some(child_info.clone()),
166 child_changes
167 .map(|(k, v)| (k.clone(), v.value().cloned()))
168 .collect(),
169 ))
170 }
171
172 self.backend.update(transaction)
173 }
174
175 pub fn commit_all(&mut self) -> Result<(), String> {
181 let changes = self.overlay.drain_storage_changes::<_, _, N>(
182 &self.backend,
183 None,
184 Default::default(),
185 &mut Default::default(),
186 )?;
187
188 self.backend.apply_transaction(changes.transaction_storage_root, changes.transaction);
189 Ok(())
190 }
191
192 pub fn execute_with<R>(&mut self, execute: impl FnOnce() -> R) -> R {
196 let mut ext = self.ext();
197 externalities::set_and_run_with_externalities(&mut ext, execute)
198 }
199
200 pub fn execute_with_safe<R>(&mut self, f: impl FnOnce() -> R + UnwindSafe) -> Result<R, String> {
205 let mut ext = AssertUnwindSafe(self.ext());
206 std::panic::catch_unwind(move ||
207 externalities::set_and_run_with_externalities(&mut *ext, f)
208 ).map_err(|e| {
209 format!("Closure panicked: {:?}", e)
210 })
211 }
212}
213
214impl<H: Hasher, N: ChangesTrieBlockNumber> std::fmt::Debug for TestExternalities<H, N>
215 where H::Out: Ord + codec::Codec,
216{
217 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
218 write!(f, "overlay: {:?}\nbackend: {:?}", self.overlay, self.backend.pairs())
219 }
220}
221
222impl<H: Hasher, N: ChangesTrieBlockNumber> PartialEq for TestExternalities<H, N>
223 where
224 H::Out: Ord + 'static + codec::Codec
225{
226 fn eq(&self, other: &TestExternalities<H, N>) -> bool {
229 self.as_backend().eq(&other.as_backend())
230 }
231}
232
233impl<H: Hasher, N: ChangesTrieBlockNumber> Default for TestExternalities<H, N>
234 where
235 H::Out: Ord + 'static + codec::Codec,
236{
237 fn default() -> Self { Self::new(Default::default()) }
238}
239
240impl<H: Hasher, N: ChangesTrieBlockNumber> From<Storage> for TestExternalities<H, N>
241 where
242 H::Out: Ord + 'static + codec::Codec,
243{
244 fn from(storage: Storage) -> Self {
245 Self::new(storage)
246 }
247}
248
249impl<H, N> externalities::ExtensionStore for TestExternalities<H, N> where
250 H: Hasher,
251 H::Out: Ord + codec::Codec,
252 N: ChangesTrieBlockNumber,
253{
254 fn extension_by_type_id(&mut self, type_id: TypeId) -> Option<&mut dyn Any> {
255 self.extensions.get_mut(type_id)
256 }
257
258 fn register_extension_with_type_id(
259 &mut self,
260 type_id: TypeId,
261 extension: Box<dyn Extension>,
262 ) -> Result<(), externalities::Error> {
263 self.extensions.register_with_type_id(type_id, extension)
264 }
265
266 fn deregister_extension_by_type_id(&mut self, type_id: TypeId) -> Result<(), externalities::Error> {
267 if self.extensions.deregister(type_id) {
268 Ok(())
269 } else {
270 Err(externalities::Error::ExtensionIsNotRegistered(type_id))
271 }
272 }
273}
274
275#[cfg(test)]
276mod tests {
277 use super::*;
278 use tet_core::{H256, traits::Externalities, storage::ChildInfo};
279 use tp_runtime::traits::BlakeTwo256;
280 use hex_literal::hex;
281
282 #[test]
283 fn commit_should_work() {
284 let mut ext = TestExternalities::<BlakeTwo256, u64>::default();
285 let mut ext = ext.ext();
286 ext.set_storage(b"doe".to_vec(), b"reindeer".to_vec());
287 ext.set_storage(b"dog".to_vec(), b"puppy".to_vec());
288 ext.set_storage(b"dogglesworth".to_vec(), b"cat".to_vec());
289 let root = H256::from(hex!("2a340d3dfd52f5992c6b117e9e45f479e6da5afffafeb26ab619cf137a95aeb8"));
290 assert_eq!(H256::from_slice(ext.storage_root().as_slice()), root);
291 }
292
293 #[test]
294 fn set_and_retrieve_code() {
295 let mut ext = TestExternalities::<BlakeTwo256, u64>::default();
296 let mut ext = ext.ext();
297
298 let code = vec![1, 2, 3];
299 ext.set_storage(CODE.to_vec(), code.clone());
300
301 assert_eq!(&ext.storage(CODE).unwrap(), &code);
302 }
303
304 #[test]
305 fn check_send() {
306 fn assert_send<T: Send>() {}
307 assert_send::<TestExternalities::<BlakeTwo256, u64>>();
308 }
309
310 #[test]
311 fn commit_all_and_kill_child_storage() {
312 let mut ext = TestExternalities::<BlakeTwo256, u64>::default();
313 let child_info = ChildInfo::new_default(&b"test_child"[..]);
314
315 {
316 let mut ext = ext.ext();
317 ext.place_child_storage(&child_info, b"doe".to_vec(), Some(b"reindeer".to_vec()));
318 ext.place_child_storage(&child_info, b"dog".to_vec(), Some(b"puppy".to_vec()));
319 ext.place_child_storage(&child_info, b"dog2".to_vec(), Some(b"puppy2".to_vec()));
320 }
321
322 ext.commit_all().unwrap();
323
324 {
325 let mut ext = ext.ext();
326
327 assert!(!ext.kill_child_storage(&child_info, Some(2)), "Should not delete all keys");
328
329 assert!(ext.child_storage(&child_info, &b"doe"[..]).is_none());
330 assert!(ext.child_storage(&child_info, &b"dog"[..]).is_none());
331 assert!(ext.child_storage(&child_info, &b"dog2"[..]).is_some());
332 }
333 }
334
335 #[test]
336 fn as_backend_generates_same_backend_as_commit_all() {
337 let mut ext = TestExternalities::<BlakeTwo256, u64>::default();
338 {
339 let mut ext = ext.ext();
340 ext.set_storage(b"doe".to_vec(), b"reindeer".to_vec());
341 ext.set_storage(b"dog".to_vec(), b"puppy".to_vec());
342 ext.set_storage(b"dogglesworth".to_vec(), b"cat".to_vec());
343 }
344
345 let backend = ext.as_backend();
346
347 ext.commit_all().unwrap();
348 assert!(ext.backend.eq(&backend), "Both backend should be equal.");
349 }
350}