@@ -42,6 +42,7 @@ use sp_runtime::{
4242} ;
4343
4444const TEST_KEY : & [ u8 ] = b":test:key:" ;
45+ const TEST_KEY_2 : & [ u8 ] = b":test:key_2:" ;
4546
4647#[ frame_support:: pallet( dev_mode) ]
4748mod custom {
@@ -354,6 +355,7 @@ impl frame_system::Config for Runtime {
354355 type PostTransactions = MockedSystemCallbacks ;
355356 type MultiBlockMigrator = MockedModeGetter ;
356357 type ExtensionsWeightInfo = MockExtensionsWeights ;
358+ type SingleBlockMigrations = CustomOnRuntimeUpgrade ;
357359}
358360
359361#[ derive(
@@ -479,17 +481,24 @@ type TestBlock = Block<UncheckedXt>;
479481// Will contain `true` when the custom runtime logic was called.
480482const CUSTOM_ON_RUNTIME_KEY : & [ u8 ] = b":custom:on_runtime" ;
481483
482- struct CustomOnRuntimeUpgrade ;
484+ pub struct CustomOnRuntimeUpgrade ;
483485impl OnRuntimeUpgrade for CustomOnRuntimeUpgrade {
484486 fn on_runtime_upgrade ( ) -> Weight {
485487 sp_io:: storage:: set ( TEST_KEY , "custom_upgrade" . as_bytes ( ) ) ;
488+ sp_io:: storage:: set ( TEST_KEY_2 , "try_runtime_upgrade_works" . as_bytes ( ) ) ;
486489 sp_io:: storage:: set ( CUSTOM_ON_RUNTIME_KEY , & true . encode ( ) ) ;
487490 System :: deposit_event ( frame_system:: Event :: CodeUpdated ) ;
488491
489492 assert_eq ! ( 0 , System :: last_runtime_upgrade_spec_version( ) ) ;
490493
491494 Weight :: from_parts ( 100 , 0 )
492495 }
496+
497+ #[ cfg( feature = "try-runtime" ) ]
498+ fn post_upgrade ( _state : Vec < u8 > ) -> Result < ( ) , TryRuntimeError > {
499+ assert_eq ! ( & sp_io:: storage:: get( TEST_KEY_2 ) . unwrap( ) [ ..] , * b"try_runtime_upgrade_works" ) ;
500+ Ok ( ( ) )
501+ }
493502}
494503
495504type Executive = super :: Executive <
@@ -1074,20 +1083,15 @@ fn all_weights_are_recorded_correctly() {
10741083 MockedSystemCallbacks :: reset ( ) ;
10751084
10761085 // All weights that show up in the `initialize_block_impl`
1077- let custom_runtime_upgrade_weight = CustomOnRuntimeUpgrade :: on_runtime_upgrade ( ) ;
1078- let runtime_upgrade_weight =
1079- <AllPalletsWithSystem as OnRuntimeUpgrade >:: on_runtime_upgrade ( ) ;
1086+ let runtime_upgrade_weight = Executive :: execute_on_runtime_upgrade ( ) ;
10801087 let on_initialize_weight =
10811088 <AllPalletsWithSystem as OnInitialize < u64 > >:: on_initialize ( block_number) ;
10821089 let base_block_weight = <Runtime as frame_system:: Config >:: BlockWeights :: get ( ) . base_block ;
10831090
10841091 // Weights are recorded correctly
10851092 assert_eq ! (
10861093 frame_system:: Pallet :: <Runtime >:: block_weight( ) . total( ) ,
1087- custom_runtime_upgrade_weight +
1088- runtime_upgrade_weight +
1089- on_initialize_weight +
1090- base_block_weight,
1094+ runtime_upgrade_weight + on_initialize_weight + base_block_weight,
10911095 ) ;
10921096 } ) ;
10931097}
@@ -1294,6 +1298,35 @@ fn try_execute_block_works() {
12941298 } ) ;
12951299}
12961300
1301+ #[ test]
1302+ #[ cfg( feature = "try-runtime" ) ]
1303+ fn try_runtime_upgrade_works ( ) {
1304+ use frame_support:: traits:: OnGenesis ;
1305+
1306+ sp_tracing:: init_for_tests ( ) ;
1307+
1308+ type ExecutiveWithoutMigrations = super :: Executive <
1309+ Runtime ,
1310+ Block < UncheckedXt > ,
1311+ ChainContext < Runtime > ,
1312+ Runtime ,
1313+ AllPalletsWithSystem ,
1314+ > ;
1315+
1316+ new_test_ext ( 1 ) . execute_with ( || {
1317+ // Call `on_genesis` to reset the storage version of all pallets.
1318+ AllPalletsWithSystem :: on_genesis ( ) ;
1319+
1320+ // Make sure the test storages are un-set
1321+ assert ! ( & sp_io:: storage:: get( TEST_KEY_2 ) . is_none( ) ) ;
1322+
1323+ ExecutiveWithoutMigrations :: try_runtime_upgrade ( UpgradeCheckSelect :: All ) . unwrap ( ) ;
1324+
1325+ // Make sure the test storages were set
1326+ assert_eq ! ( & sp_io:: storage:: get( TEST_KEY_2 ) . unwrap( ) [ ..] , * b"try_runtime_upgrade_works" ) ;
1327+ } ) ;
1328+ }
1329+
12971330/// Same as `extrinsic_while_exts_forbidden_errors` but using the try-runtime function.
12981331#[ test]
12991332#[ cfg( feature = "try-runtime" ) ]
0 commit comments