1
// Copyright 2024 Moonbeam foundation
2
// This file is part of Moonbeam.
3

            
4
// Moonbeam is free software: you can redistribute it and/or modify
5
// it under the terms of the GNU General Public License as published by
6
// the Free Software Foundation, either version 3 of the License, or
7
// (at your option) any later version.
8

            
9
// Moonbeam is distributed in the hope that it will be useful,
10
// but WITHOUT ANY WARRANTY; without even the implied warranty of
11
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12
// GNU General Public License for more details.
13

            
14
// You should have received a copy of the GNU General Public License
15
// along with Moonbeam.  If not, see <http://www.gnu.org/licenses/>.
16

            
17
//! Unit testing
18
use {
19
	crate::{
20
		mock::{ExtBuilder, LazyMigrations, RuntimeOrigin, Test},
21
		Error, StateMigrationStatus, StateMigrationStatusValue, MAX_ITEM_PROOF_SIZE,
22
		PROOF_SIZE_BUFFER,
23
	},
24
	frame_support::{assert_noop, assert_ok, traits::Hooks, weights::Weight},
25
	rlp::RlpStream,
26
	sp_core::{H160, H256},
27
	sp_io::hashing::keccak_256,
28
	sp_runtime::{traits::Bounded, AccountId32},
29
};
30

            
31
use pallet_evm::AddressMapping;
32

            
33
// Helper function that calculates the contract address
34
113
pub fn contract_address(sender: H160, nonce: u64) -> H160 {
35
113
	let mut rlp = RlpStream::new_list(2);
36
113
	rlp.append(&sender);
37
113
	rlp.append(&nonce);
38
113

            
39
113
	H160::from_slice(&keccak_256(&rlp.out())[12..])
40
113
}
41

            
42
115
fn address_build(seed: u8) -> H160 {
43
115
	let address = H160::from(H256::from(keccak_256(&[seed; 32])));
44
115
	address
45
115
}
46

            
47
// Helper function that creates a `num_entries` storage entries for a contract
48
113
fn mock_contract_with_entries(seed: u8, nonce: u64, num_entries: u32) -> H160 {
49
113
	let address = address_build(seed);
50
113

            
51
113
	let contract_address = contract_address(address, nonce);
52
113
	let account_id =
53
113
		<Test as pallet_evm::Config>::AddressMapping::into_account_id(contract_address);
54
113
	let _ = frame_system::Pallet::<Test>::inc_sufficients(&account_id);
55

            
56
	// Add num_entries storage entries to the suicided contract
57
12141
	for i in 0..num_entries {
58
12141
		pallet_evm::AccountStorages::<Test>::insert(
59
12141
			contract_address,
60
12141
			H256::from_low_u64_be(i as u64),
61
12141
			H256::from_low_u64_be(i as u64),
62
12141
		);
63
12141
	}
64

            
65
113
	contract_address
66
113
}
67

            
68
1
fn create_dummy_contract_without_metadata(seed: u8) -> H160 {
69
1
	let address = address_build(seed);
70
1
	let dummy_code = vec![1, 2, 3];
71
1
	pallet_evm::AccountCodes::<Test>::insert(address, dummy_code);
72
1
	address
73
1
}
74

            
75
#[test]
76
1
fn test_clear_suicided_contract_succesfull() {
77
1
	ExtBuilder::default().build().execute_with(|| {
78
1
		let contract_address = mock_contract_with_entries(1, 1, 10);
79
1

            
80
1
		// No addresses have been migrated yet
81
1
		assert_eq!(crate::pallet::SuicidedContractsRemoved::<Test>::get(), 0);
82

            
83
		// The account has some storage entries
84
1
		assert_eq!(
85
1
			pallet_evm::AccountStorages::<Test>::iter_prefix(contract_address).count(),
86
1
			10
87
1
		);
88

            
89
		// Call the extrinsic to delete the storage entries
90
1
		let _ = LazyMigrations::clear_suicided_storage(
91
1
			RuntimeOrigin::signed(AccountId32::from([45; 32])),
92
1
			vec![contract_address].try_into().unwrap(),
93
1
			1000,
94
1
		);
95
1

            
96
1
		// One address has been migrated
97
1
		assert_eq!(crate::pallet::SuicidedContractsRemoved::<Test>::get(), 1);
98
		// All the account storage should have been removed
99
1
		assert_eq!(
100
1
			pallet_evm::AccountStorages::<Test>::iter_prefix(contract_address).count(),
101
1
			0
102
1
		);
103
1
	})
104
1
}
105

            
106
// Test that the extrinsic fails if the contract is not suicided
107
#[test]
108
1
fn test_clear_suicided_contract_failed() {
109
1
	ExtBuilder::default().build().execute_with(|| {
110
1
		let contract1_address = mock_contract_with_entries(1, 1, 10);
111
1
		let contract2_address = mock_contract_with_entries(2, 1, 10);
112
1

            
113
1
		// The contracts have not been self-destructed.
114
1
		pallet_evm::AccountCodes::<Test>::insert(contract1_address, vec![1, 2, 3]);
115
1
		pallet_evm::Suicided::<Test>::insert(contract2_address, ());
116
1

            
117
1
		assert_noop!(
118
1
			LazyMigrations::clear_suicided_storage(
119
1
				RuntimeOrigin::signed(AccountId32::from([45; 32])),
120
1
				vec![contract1_address].try_into().unwrap(),
121
1
				1000
122
1
			),
123
1
			Error::<Test>::ContractNotCorrupted
124
1
		);
125

            
126
1
		assert_noop!(
127
1
			LazyMigrations::clear_suicided_storage(
128
1
				RuntimeOrigin::signed(AccountId32::from([45; 32])),
129
1
				vec![contract2_address].try_into().unwrap(),
130
1
				1000
131
1
			),
132
1
			Error::<Test>::ContractNotCorrupted
133
1
		);
134

            
135
		// Check that no storage has been removed
136

            
137
1
		assert_eq!(
138
1
			pallet_evm::AccountStorages::<Test>::iter_prefix(contract1_address).count(),
139
1
			10
140
1
		);
141
1
		assert_eq!(
142
1
			pallet_evm::AccountStorages::<Test>::iter_prefix(contract2_address).count(),
143
1
			10
144
1
		);
145
1
	})
146
1
}
147

            
148
// Test that the extrinsic can handle an empty input
149
#[test]
150
1
fn test_clear_suicided_empty_input() {
151
1
	ExtBuilder::default().build().execute_with(|| {
152
1
		let contract_address = mock_contract_with_entries(1, 1, 10);
153
1

            
154
1
		let _ = LazyMigrations::clear_suicided_storage(
155
1
			RuntimeOrigin::signed(AccountId32::from([45; 32])),
156
1
			vec![].try_into().unwrap(),
157
1
			1000,
158
1
		);
159
1

            
160
1
		assert_eq!(
161
1
			pallet_evm::AccountStorages::<Test>::iter_prefix(contract_address).count(),
162
1
			10
163
1
		);
164
1
	})
165
1
}
166

            
167
// Test with multiple deleted contracts ensuring that the extrinsic can handle
168
// multiple addresses at once.
169
#[test]
170
1
fn test_clear_suicided_contract_multiple_addresses() {
171
1
	ExtBuilder::default().build().execute_with(|| {
172
1
		let contract_address1 = mock_contract_with_entries(1, 1, 10);
173
1
		let contract_address2 = mock_contract_with_entries(2, 1, 20);
174
1
		let contract_address3 = mock_contract_with_entries(3, 1, 30);
175
1

            
176
1
		// Call the extrinsic to delete the storage entries
177
1
		let _ = LazyMigrations::clear_suicided_storage(
178
1
			RuntimeOrigin::signed(AccountId32::from([45; 32])),
179
1
			vec![contract_address1, contract_address2, contract_address3]
180
1
				.try_into()
181
1
				.unwrap(),
182
1
			1000,
183
1
		)
184
1
		.unwrap();
185
1

            
186
1
		assert_eq!(
187
1
			pallet_evm::AccountStorages::<Test>::iter_prefix(contract_address1).count(),
188
1
			0
189
1
		);
190
1
		assert_eq!(
191
1
			pallet_evm::AccountStorages::<Test>::iter_prefix(contract_address2).count(),
192
1
			0
193
1
		);
194
1
		assert_eq!(
195
1
			pallet_evm::AccountStorages::<Test>::iter_prefix(contract_address3).count(),
196
1
			0
197
1
		);
198
1
	})
199
1
}
200

            
201
// Test that the limit of entries to be deleted is respected
202
#[test]
203
1
fn test_clear_suicided_entry_limit() {
204
1
	ExtBuilder::default().build().execute_with(|| {
205
1
		let contract_address1 = mock_contract_with_entries(1, 1, 2000);
206
1
		let contract_address2 = mock_contract_with_entries(2, 1, 1);
207
1

            
208
1
		let _ = LazyMigrations::clear_suicided_storage(
209
1
			RuntimeOrigin::signed(AccountId32::from([45; 32])),
210
1
			vec![contract_address1, contract_address2]
211
1
				.try_into()
212
1
				.unwrap(),
213
1
			1000,
214
1
		)
215
1
		.unwrap();
216
1
		assert_eq!(
217
1
			pallet_evm::AccountStorages::<Test>::iter_prefix(contract_address1).count(),
218
1
			1000
219
1
		);
220

            
221
1
		assert_eq!(
222
1
			pallet_evm::AccountStorages::<Test>::iter_prefix(contract_address2).count(),
223
1
			1
224
1
		);
225
1
	})
226
1
}
227

            
228
// Test a combination of Suicided and non-suicided contracts
229
#[test]
230
1
fn test_clear_suicided_mixed_suicided_and_non_suicided() {
231
1
	ExtBuilder::default().build().execute_with(|| {
232
1
		let contract_address1 = mock_contract_with_entries(1, 1, 10);
233
1
		let contract_address2 = mock_contract_with_entries(2, 1, 10);
234
1
		let contract_address3 = mock_contract_with_entries(3, 1, 10);
235
1
		let contract_address4 = mock_contract_with_entries(4, 1, 10);
236
1

            
237
1
		// Contract has not been self-destructed.
238
1
		pallet_evm::AccountCodes::<Test>::insert(contract_address3, vec![1, 2, 3]);
239
1

            
240
1
		assert_noop!(
241
1
			LazyMigrations::clear_suicided_storage(
242
1
				RuntimeOrigin::signed(AccountId32::from([45; 32])),
243
1
				vec![
244
1
					contract_address1,
245
1
					contract_address2,
246
1
					contract_address3,
247
1
					contract_address4
248
1
				]
249
1
				.try_into()
250
1
				.unwrap(),
251
1
				1000
252
1
			),
253
1
			Error::<Test>::ContractNotCorrupted
254
1
		);
255

            
256
1
		assert_eq!(
257
1
			pallet_evm::AccountStorages::<Test>::iter_prefix(contract_address1).count(),
258
1
			10
259
1
		);
260
1
		assert_eq!(
261
1
			pallet_evm::AccountStorages::<Test>::iter_prefix(contract_address2).count(),
262
1
			10
263
1
		);
264
1
		assert_eq!(
265
1
			pallet_evm::AccountStorages::<Test>::iter_prefix(contract_address3).count(),
266
1
			10
267
1
		);
268
1
		assert_eq!(
269
1
			pallet_evm::AccountStorages::<Test>::iter_prefix(contract_address4).count(),
270
1
			10
271
1
		);
272
1
	})
273
1
}
274

            
275
#[test]
276
1
fn test_create_contract_metadata_contract_not_exist() {
277
1
	ExtBuilder::default().build().execute_with(|| {
278
1
		assert_noop!(
279
1
			LazyMigrations::create_contract_metadata(
280
1
				RuntimeOrigin::signed(AccountId32::from([45; 32])),
281
1
				address_build(1),
282
1
			),
283
1
			Error::<Test>::ContractNotExist
284
1
		);
285
1
	});
286
1
}
287

            
288
#[test]
289
1
fn test_create_contract_metadata_success_path() {
290
1
	ExtBuilder::default().build().execute_with(|| {
291
1
		// Setup: create a dummy contract
292
1
		let address = create_dummy_contract_without_metadata(1);
293
1

            
294
1
		assert_ok!(LazyMigrations::create_contract_metadata(
295
1
			RuntimeOrigin::signed(AccountId32::from([45; 32])),
296
1
			address,
297
1
		));
298

            
299
1
		assert!(pallet_evm::AccountCodesMetadata::<Test>::get(address).is_some());
300

            
301
		// Should not be able to set metadata again
302
1
		assert_noop!(
303
1
			LazyMigrations::create_contract_metadata(
304
1
				RuntimeOrigin::signed(AccountId32::from([45; 32])),
305
1
				address,
306
1
			),
307
1
			Error::<Test>::ContractMetadataAlreadySet
308
1
		);
309
1
	});
310
1
}
311

            
312
4
fn count_keys_and_data_without_code() -> (u64, u64) {
313
4
	let mut keys: u64 = 0;
314
4
	let mut data: u64 = 0;
315
4

            
316
4
	let mut current_key: Option<Vec<u8>> = Some(Default::default());
317
10148
	while let Some(key) = current_key {
318
10144
		if key.as_slice() == sp_core::storage::well_known_keys::CODE {
319
4
			current_key = sp_io::storage::next_key(&key);
320
4
			continue;
321
10140
		}
322
10140
		keys += 1;
323
10140
		if let Some(_) = sp_io::storage::get(&key) {
324
10136
			data += 1;
325
10136
		}
326
10140
		current_key = sp_io::storage::next_key(&key);
327
	}
328

            
329
4
	(keys, data)
330
4
}
331

            
332
210
fn weight_for(read: u64, write: u64) -> Weight {
333
210
	<Test as frame_system::Config>::DbWeight::get().reads_writes(read, write)
334
210
}
335

            
336
206
fn rem_weight_for_entries(num_entries: u64) -> Weight {
337
206
	let proof = PROOF_SIZE_BUFFER + num_entries * MAX_ITEM_PROOF_SIZE;
338
206
	Weight::from_parts(u64::max_value(), proof)
339
206
}
340

            
341
#[test]
342
1
fn test_state_migration_baseline() {
343
1
	ExtBuilder::default().build().execute_with(|| {
344
1
		assert_eq!(
345
1
			StateMigrationStatusValue::<Test>::get(),
346
1
			(StateMigrationStatus::NotStarted, 0)
347
1
		);
348

            
349
1
		let (keys, data) = count_keys_and_data_without_code();
350
1
		println!("Keys: {}, Data: {}", keys, data);
351
1

            
352
1
		let weight = LazyMigrations::on_idle(0, Weight::max_value());
353
1

            
354
1
		// READS: 2 * keys + 2 (skipped and status)
355
1
		// Next key requests = keys (we have first key as default which is not counted, and extra
356
1
		// next_key request to check if we are done)
357
1
		//
358
1
		// 1 next key request for the skipped key ":code"
359
1
		// Read requests = keys (we read each key once)
360
1
		// 1 Read request for the StateMigrationStatusValue
361
1

            
362
1
		// WRITES: data + 1 (status)
363
1
		// Write requests = data (we write each data once)
364
1
		// 1 Write request for the StateMigrationStatusValue
365
1
		assert_eq!(weight, weight_for(2 * keys + 2, data + 1));
366

            
367
1
		assert_eq!(
368
1
			StateMigrationStatusValue::<Test>::get(),
369
1
			(StateMigrationStatus::Complete, keys)
370
1
		);
371
1
	})
372
1
}
373

            
374
#[test]
375
1
fn test_state_migration_cannot_fit_any_item() {
376
1
	ExtBuilder::default().build().execute_with(|| {
377
1
		StateMigrationStatusValue::<Test>::put((StateMigrationStatus::NotStarted, 0));
378
1

            
379
1
		let weight = LazyMigrations::on_idle(0, rem_weight_for_entries(0));
380
1

            
381
1
		assert_eq!(weight, weight_for(0, 0));
382
1
	})
383
1
}
384

            
385
#[test]
386
1
fn test_state_migration_when_complete() {
387
1
	ExtBuilder::default().build().execute_with(|| {
388
1
		StateMigrationStatusValue::<Test>::put((StateMigrationStatus::Complete, 0));
389
1

            
390
1
		let weight = LazyMigrations::on_idle(0, Weight::max_value());
391
1

            
392
1
		// just reading the status of the migration
393
1
		assert_eq!(weight, weight_for(1, 0));
394
1
	})
395
1
}
396

            
397
#[test]
398
1
fn test_state_migration_when_errored() {
399
1
	ExtBuilder::default().build().execute_with(|| {
400
1
		StateMigrationStatusValue::<Test>::put((
401
1
			StateMigrationStatus::Error("Error".as_bytes().to_vec().try_into().unwrap_or_default()),
402
1
			1,
403
1
		));
404
1

            
405
1
		let weight = LazyMigrations::on_idle(0, Weight::max_value());
406
1

            
407
1
		// just reading the status of the migration
408
1
		assert_eq!(weight, weight_for(1, 0));
409
1
	})
410
1
}
411

            
412
#[test]
413
1
fn test_state_migration_can_only_fit_one_item() {
414
1
	ExtBuilder::default().build().execute_with(|| {
415
1
		assert_eq!(
416
1
			StateMigrationStatusValue::<Test>::get(),
417
1
			(StateMigrationStatus::NotStarted, 0)
418
1
		);
419

            
420
1
		let data = sp_io::storage::get(Default::default());
421
1
		let weight = LazyMigrations::on_idle(0, rem_weight_for_entries(1));
422
1

            
423
1
		let reads = 2; // key read + status read
424
1
		let writes = 1 + data.map(|_| 1).unwrap_or(0);
425
1
		assert_eq!(weight, weight_for(reads, writes));
426

            
427
1
		assert!(matches!(
428
1
			StateMigrationStatusValue::<Test>::get(),
429
			(StateMigrationStatus::Started(_), 1)
430
		));
431

            
432
1
		let weight = LazyMigrations::on_idle(0, rem_weight_for_entries(3));
433
1
		let reads = 3 + 3 + 1; // next key + key read + status
434
1
		let writes = 1 + 3; // status write + key write
435
1
		assert_eq!(weight, weight_for(reads, writes));
436
1
	})
437
1
}
438

            
439
#[test]
440
1
fn test_state_migration_can_only_fit_three_item() {
441
1
	ExtBuilder::default().build().execute_with(|| {
442
1
		assert_eq!(
443
1
			StateMigrationStatusValue::<Test>::get(),
444
1
			(StateMigrationStatus::NotStarted, 0)
445
1
		);
446

            
447
1
		let weight = LazyMigrations::on_idle(0, rem_weight_for_entries(3));
448
1

            
449
1
		// 2 next key requests (default key dons't need a next key request) + 1 status read
450
1
		// 3 key reads.
451
1
		// 1 status write + 2 key writes (default key doesn't have any data)
452
1
		let reads = 6;
453
1
		let writes = 3;
454
1
		assert_eq!(weight, weight_for(reads, writes));
455

            
456
1
		assert!(matches!(
457
1
			StateMigrationStatusValue::<Test>::get(),
458
			(StateMigrationStatus::Started(_), 3)
459
		));
460
1
	})
461
1
}
462

            
463
#[test]
464
1
fn test_state_migration_can_fit_exactly_all_item() {
465
1
	ExtBuilder::default().build().execute_with(|| {
466
1
		assert_eq!(
467
1
			StateMigrationStatusValue::<Test>::get(),
468
1
			(StateMigrationStatus::NotStarted, 0)
469
1
		);
470

            
471
1
		let (keys, data) = count_keys_and_data_without_code();
472
1
		let weight = LazyMigrations::on_idle(0, rem_weight_for_entries(keys));
473
1

            
474
1
		// we deduct the extra next_key request to check if we are done.
475
1
		// will know if we are done on the next call to on_idle
476
1
		assert_eq!(weight, weight_for(2 * keys + 1, data + 1));
477

            
478
1
		assert!(matches!(
479
1
			StateMigrationStatusValue::<Test>::get(),
480
1
			(StateMigrationStatus::Started(_), n) if n == keys,
481
		));
482

            
483
		// after calling on_idle status is added to the storage so we need to account for that
484
1
		let (new_keys, new_data) = count_keys_and_data_without_code();
485
1
		let (diff_keys, diff_data) = (new_keys - keys, new_data - data);
486
1

            
487
1
		let weight = LazyMigrations::on_idle(0, rem_weight_for_entries(1 + diff_keys));
488
1
		// (next_key + read) for each new key + status + next_key to check if we are done
489
1
		let reads = diff_keys * 2 + 2;
490
1
		let writes = 1 + diff_data; // status
491
1
		assert_eq!(weight, weight_for(reads, writes));
492

            
493
1
		assert!(matches!(
494
1
			StateMigrationStatusValue::<Test>::get(),
495
1
			(StateMigrationStatus::Complete, n) if n == new_keys,
496
		));
497
1
	})
498
1
}
499

            
500
#[test]
501
1
fn test_state_migration_will_migrate_10_000_items() {
502
1
	ExtBuilder::default().build().execute_with(|| {
503
1
		assert_eq!(
504
1
			StateMigrationStatusValue::<Test>::get(),
505
1
			(StateMigrationStatus::NotStarted, 0)
506
1
		);
507

            
508
101
		for i in 0..100 {
509
100
			mock_contract_with_entries(i as u8, i as u64, 100);
510
100
		}
511

            
512
1
		StateMigrationStatusValue::<Test>::put((StateMigrationStatus::NotStarted, 0));
513
1

            
514
1
		let (keys, data) = count_keys_and_data_without_code();
515
1

            
516
1
		// assuming we can only fit 100 items at a time
517
1

            
518
1
		let mut total_weight: Weight = Weight::zero();
519
1
		let num_of_on_idle_calls = 200;
520
1
		let entries_per_on_idle = 100;
521
1
		let needed_on_idle_calls = (keys as f64 / entries_per_on_idle as f64).ceil() as u64;
522
1

            
523
1
		// Reads:
524
1
		// Read status => num_of_on_idle_calls
525
1
		// Read keys   => keys
526
1
		// Next keys   => keys - 1  + 1 skip + 1 done check
527
1
		//
528
1
		// Writes:
529
1
		// Write status => needed_on_idle_calls
530
1
		// Write keys   => data
531
1
		let expected_reads = (keys - 1 + 2) + keys + num_of_on_idle_calls;
532
1
		let expected_writes = data + needed_on_idle_calls;
533
1

            
534
1
		println!("Keys: {}, Data: {}", keys, data);
535
1
		println!("entries_per_on_idle: {}", entries_per_on_idle);
536
1
		println!("num_of_on_idle_calls: {}", num_of_on_idle_calls);
537
1
		println!("needed_on_idle_calls: {}", needed_on_idle_calls);
538
1
		println!(
539
1
			"Expected Reads: {}, Expected Writes: {}",
540
1
			expected_reads, expected_writes
541
1
		);
542

            
543
200
		for i in 1..=num_of_on_idle_calls {
544
200
			let weight = LazyMigrations::on_idle(i, rem_weight_for_entries(entries_per_on_idle));
545
200
			total_weight = total_weight.saturating_add(weight);
546
200

            
547
200
			let status = StateMigrationStatusValue::<Test>::get();
548
200
			if i < needed_on_idle_calls {
549
101
				let migrated_so_far = i * entries_per_on_idle;
550
101
				assert!(
551
101
					matches!(status, (StateMigrationStatus::Started(_), n) if n == migrated_so_far),
552
					"Status: {:?} at call: #{} doesn't match Started",
553
					status,
554
					i,
555
				);
556
101
				assert!(weight.all_gte(weight_for(1, 0)));
557
			} else {
558
99
				assert!(
559
99
					matches!(status, (StateMigrationStatus::Complete, n) if n == keys),
560
					"Status: {:?} at call: {} doesn't match Complete",
561
					status,
562
					i,
563
				);
564
99
				if i == needed_on_idle_calls {
565
					// last call to on_idle
566
1
					assert!(weight.all_gte(weight_for(1, 0)));
567
				} else {
568
					// extra calls to on_idle, just status update check
569
98
					assert_eq!(weight, weight_for(1, 0));
570
				}
571
			}
572
		}
573

            
574
1
		assert_eq!(total_weight, weight_for(expected_reads, expected_writes));
575
1
	})
576
1
}