Rework the runtime upgrade test (#727)

* Rework the runtime upgrade test

* Update test/service/tests/runtime_upgrade.rs

* Update test/service/tests/runtime_upgrade.rs

* Update Cargo.lock

* FMT

Co-authored-by: Shawn Tabrizi <shawntabrizi@gmail.com>

cumulus/Cargo.toml

@@ -35,7 +35,6 @@ members = [
 	"test/relay-sproof-builder",
 	"test/relay-validation-worker-provider",
 	"test/runtime",
-	"test/runtime-upgrade",
 	"test/service",
 ]
 

cumulus/test/runtime-upgrade/build.rs

@@ -1 +0,0 @@
-../runtime/build.rs

cumulus/test/runtime-upgrade/src

@@ -1 +0,0 @@
-../runtime/src

cumulus/test/runtime/Cargo.toml

@@ -67,4 +67,4 @@ std = [
 	"sp-transaction-pool/std",
 	"sp-version/std",
 ]
-upgrade = []
+increment-spec-version = []

cumulus/test/runtime/build.rs

@@ -21,5 +21,12 @@ fn main() {
 		.with_current_project()
 		.export_heap_base()
 		.import_memory()
-		.build()
+		.build();
+
+	WasmBuilder::new()
+		.with_current_project()
+		.enable_feature("increment-spec-version")
+		.import_memory()
+		.set_file_name("wasm_binary_spec_version_incremented.rs")
+		.build();
 }

cumulus/test/runtime/src/lib.rs

@@ -22,6 +22,12 @@
 #[cfg(feature = "std")]
 include!(concat!(env!("OUT_DIR"), "/wasm_binary.rs"));
 
+pub mod wasm_spec_version_incremented {
+	#[cfg(feature = "std")]
+	include!(concat!(env!("OUT_DIR"), "/wasm_binary_spec_version_incremented.rs"));
+}
+
+use frame_support::traits::OnRuntimeUpgrade;
 use sp_api::{decl_runtime_apis, impl_runtime_apis};
 use sp_core::OpaqueMetadata;
 use sp_runtime::{
@@ -59,9 +65,13 @@ impl_opaque_keys! {
 	pub struct SessionKeys {}
 }
 
+/// Some key that we set in genesis and only read in [`TestRuntimeUpgrade`] to ensure that
+/// [`OnRuntimeUpgrade`] works as expected.
+pub const TEST_RUNTIME_UPGRADE_KEY: &[u8] = b"+test_runtime_upgrade_key+";
+
 // The only difference between the two declarations below is the `spec_version`. With the
-// `upgrade` feature enabled `spec_version` should be greater than the one of without the
-// `upgrade` feature.
+// `increment-spec-version` feature enabled `spec_version` should be greater than the one of without the
+// `increment-spec-version` feature.
 //
 // The duplication here is unfortunate necessity.
 //
@@ -70,27 +80,27 @@ impl_opaque_keys! {
 // details. Since macro kicks in early, it operates on AST. Thus you cannot use constants.
 // Macros are expanded top to bottom, meaning we also cannot use `cfg` here.
 
-#[cfg(feature = "upgrade")]
+#[cfg(not(feature = "increment-spec-version"))]
 #[sp_version::runtime_version]
 pub const VERSION: RuntimeVersion = RuntimeVersion {
 	spec_name: create_runtime_str!("cumulus-test-parachain"),
 	impl_name: create_runtime_str!("cumulus-test-parachain"),
 	authoring_version: 1,
 	// Read the note above.
-	spec_version: 4,
+	spec_version: 1,
 	impl_version: 1,
 	apis: RUNTIME_API_VERSIONS,
 	transaction_version: 1,
 };
 
-#[cfg(not(feature = "upgrade"))]
+#[cfg(feature = "increment-spec-version")]
 #[sp_version::runtime_version]
 pub const VERSION: RuntimeVersion = RuntimeVersion {
 	spec_name: create_runtime_str!("cumulus-test-parachain"),
 	impl_name: create_runtime_str!("cumulus-test-parachain"),
 	authoring_version: 1,
 	// Read the note above.
-	spec_version: 3,
+	spec_version: 2,
 	impl_version: 1,
 	apis: RUNTIME_API_VERSIONS,
 	transaction_version: 1,
@@ -321,10 +331,21 @@ pub type Executive = frame_executive::Executive<
 	frame_system::ChainContext<Runtime>,
 	Runtime,
 	AllPallets,
+	TestOnRuntimeUpgrade,
 >;
 /// The payload being signed in transactions.
 pub type SignedPayload = generic::SignedPayload<Call, SignedExtra>;
 
+pub struct TestOnRuntimeUpgrade;
+
+impl OnRuntimeUpgrade for TestOnRuntimeUpgrade {
+	fn on_runtime_upgrade() -> frame_support::weights::Weight {
+		assert_eq!(sp_io::storage::get(TEST_RUNTIME_UPGRADE_KEY), Some(vec![1, 2, 3, 4]));
+
+		1
+	}
+}
+
 decl_runtime_apis! {
 	pub trait GetLastTimestamp {
 		/// Returns the last timestamp of a runtime.

cumulus/test/service/Cargo.toml

@@ -27,6 +27,7 @@ sc-transaction-pool = { git = "https://github.com/paritytech/substrate", branch
 sp-arithmetic = { git = "https://github.com/paritytech/substrate", branch = "master" }
 sp-blockchain = { git = "https://github.com/paritytech/substrate", branch = "master" }
 sp-core = { git = "https://github.com/paritytech/substrate", branch = "master" }
+sp-io = { git = "https://github.com/paritytech/substrate", branch = "master" }
 sp-keyring = { git = "https://github.com/paritytech/substrate", branch = "master" }
 sp-runtime = { git = "https://github.com/paritytech/substrate", default-features = false, branch = "master" }
 sp-state-machine = { git = "https://github.com/paritytech/substrate", branch = "master" }
@@ -59,10 +60,5 @@ futures = "0.3.5"
 polkadot-test-service = { git = "https://github.com/paritytech/polkadot", branch = "master" }
 
 # Substrate dependencies
-sc-cli = { git = "https://github.com/paritytech/substrate", branch = "master" }
-substrate-test-utils = { git = "https://github.com/paritytech/substrate", branch = "master" }
-sp-maybe-compressed-blob = { git = "https://github.com/paritytech/substrate", branch = "master" }
-sp-version = { git = "https://github.com/paritytech/substrate", branch = "master" }
-
-# Cumulus
-cumulus-test-runtime-upgrade = { path = "../runtime-upgrade" }
+sc-cli = { git = "https://github.com/paritytech/substrate", branch = "polkadot-v0.9.12" }
+substrate-test-utils = { git = "https://github.com/paritytech/substrate", branch = "polkadot-v0.9.12" }

cumulus/test/service/src/chain_spec.rs

@@ -25,7 +25,24 @@ use sp_core::{sr25519, Pair, Public};
 use sp_runtime::traits::{IdentifyAccount, Verify};
 
 /// Specialized `ChainSpec` for the normal parachain runtime.
-pub type ChainSpec = sc_service::GenericChainSpec<cumulus_test_runtime::GenesisConfig, Extensions>;
+pub type ChainSpec = sc_service::GenericChainSpec<GenesisExt, Extensions>;
+
+/// Extension for the genesis config to add custom keys easily.
+#[derive(serde::Serialize, serde::Deserialize)]
+pub struct GenesisExt {
+	/// The runtime genesis config.
+	runtime_genesis_config: cumulus_test_runtime::GenesisConfig,
+}
+
+impl sp_runtime::BuildStorage for GenesisExt {
+	fn assimilate_storage(&self, storage: &mut sp_core::storage::Storage) -> Result<(), String> {
+		sp_state_machine::BasicExternalities::execute_with_storage(storage, || {
+			sp_io::storage::set(cumulus_test_runtime::TEST_RUNTIME_UPGRADE_KEY, &vec![1, 2, 3, 4]);
+		});
+
+		self.runtime_genesis_config.assimilate_storage(storage)
+	}
+}
 
 /// Helper function to generate a crypto pair from seed
 pub fn get_from_seed<TPublic: Public>(seed: &str) -> <TPublic::Pair as Pair>::Public {
@@ -65,7 +82,7 @@ pub fn get_chain_spec(id: ParaId) -> ChainSpec {
 		"Local Testnet",
 		"local_testnet",
 		ChainType::Local,
-		move || local_testnet_genesis(),
+		move || GenesisExt { runtime_genesis_config: local_testnet_genesis() },
 		vec![],
 		None,
 		None,

cumulus/test/service/tests/runtime_upgrade.rs

@@ -73,17 +73,11 @@ async fn test_runtime_upgrade() {
 		.expect("Runtime version exists");
 	expected_runtime_version.spec_version += 1;
 
-	// Replace the runtime version in the WASM blob to make it look like a new runtime.
-	let wasm = sp_maybe_compressed_blob::decompress(
-		cumulus_test_runtime_upgrade::WASM_BINARY.unwrap(),
-		sp_maybe_compressed_blob::CODE_BLOB_BOMB_LIMIT,
-	)
-	.expect("Decompressing the WASM blob works");
-	let wasm = sp_version::embed::embed_runtime_version(&wasm, expected_runtime_version.clone())
-		.expect("Embedding the runtime version works");
+	let wasm = cumulus_test_runtime::wasm_spec_version_incremented::WASM_BINARY
+		.expect("Wasm binary with incremented spec version should have been built");
 
 	// schedule runtime upgrade
-	charlie.schedule_upgrade(wasm).await.unwrap();
+	charlie.schedule_upgrade(wasm.into()).await.unwrap();
 
 	let mut import_stream = dave.client.import_notification_stream();