From 59201a40e9afe128a3869ef9ae9cb92e338551b2 Mon Sep 17 00:00:00 2001
From: Chris Sosnin <chris125_@live.com>
Date: Tue, 31 Aug 2021 18:50:39 +0300
Subject: [PATCH] Update tests

---
 pallets/parachain-system/src/tests.rs | 44 +++++++++++++++++++++++++--
 test/relay-sproof-builder/src/lib.rs  | 31 ++++++++++++++++++-
 2 files changed, 71 insertions(+), 4 deletions(-)

diff --git a/pallets/parachain-system/src/tests.rs b/pallets/parachain-system/src/tests.rs
index 3108f870a93..fc050cd3b82 100755
--- a/pallets/parachain-system/src/tests.rs
+++ b/pallets/parachain-system/src/tests.rs
@@ -406,8 +406,10 @@ fn block_tests_run_on_drop() {
 #[test]
 fn events() {
 	BlockTests::new()
-		.with_relay_sproof_builder(|_, _, builder| {
-			builder.host_config.validation_upgrade_delay = 1000;
+		.with_relay_sproof_builder(|_, block_number, builder| {
+			if block_number > 123 {
+				builder.upgrade_go_ahead = Some(relay_chain::v1::UpgradeGoAhead::GoAhead);
+			}
 		})
 		.add_with_post_test(
 			123,
@@ -421,7 +423,7 @@ fn events() {
 				let events = System::events();
 				assert_eq!(
 					events[0].event,
-					Event::ParachainSystem(crate::Event::ValidationFunctionStored(1123).into())
+					Event::ParachainSystem(crate::Event::ValidationFunctionStored.into())
 				);
 			},
 		)
@@ -461,6 +463,11 @@ fn non_overlapping() {
 #[test]
 fn manipulates_storage() {
 	BlockTests::new()
+		.with_relay_sproof_builder(|_, block_number, builder| {
+			if block_number > 123 {
+				builder.upgrade_go_ahead = Some(relay_chain::v1::UpgradeGoAhead::GoAhead);
+			}
+		})
 		.add(123, || {
 			assert!(
 				!<PendingValidationCode<Test>>::exists(),
@@ -487,6 +494,37 @@ fn manipulates_storage() {
 		);
 }
 
+#[test]
+fn aborted_upgrade() {
+	BlockTests::new()
+		.with_relay_sproof_builder(|_, block_number, builder| {
+			if block_number > 123 {
+				builder.upgrade_go_ahead = Some(relay_chain::v1::UpgradeGoAhead::Abort);
+			}
+		})
+		.add(123, || {
+			assert_ok!(System::set_code(
+				RawOrigin::Root.into(),
+				Default::default()
+			));
+		})
+		.add_with_post_test(
+			1234,
+			|| {},
+			|| {
+				assert!(
+					!<PendingValidationCode<Test>>::exists(),
+					"validation function must have been unset"
+				);
+				let events = System::events();
+				assert_eq!(
+					events[0].event,
+					Event::ParachainSystem(crate::Event::ValidationFunctionDiscarded.into())
+				);
+			},
+		);
+}
+
 #[test]
 fn checks_size() {
 	BlockTests::new()
diff --git a/test/relay-sproof-builder/src/lib.rs b/test/relay-sproof-builder/src/lib.rs
index 2c7ad67ba5e..20aa4b3559a 100644
--- a/test/relay-sproof-builder/src/lib.rs
+++ b/test/relay-sproof-builder/src/lib.rs
@@ -17,12 +17,13 @@
 use cumulus_primitives_core::{
 	relay_chain, AbridgedHostConfiguration, AbridgedHrmpChannel, ParaId,
 };
+use polkadot_primitives::v1::UpgradeGoAhead;
 use sp_runtime::traits::HashFor;
 use sp_state_machine::MemoryDB;
 use sp_std::collections::btree_map::BTreeMap;
 
 /// Builds a sproof (portmanteau of 'spoof' and 'proof') of the relay chain state.
-#[derive(Clone)]
+// #[derive(Clone)]
 pub struct RelayStateSproofBuilder {
 	/// The para id of the current parachain.
 	///
@@ -36,6 +37,7 @@ pub struct RelayStateSproofBuilder {
 
 	pub host_config: AbridgedHostConfiguration,
 	pub dmq_mqc_head: Option<relay_chain::Hash>,
+	pub upgrade_go_ahead: Option<UpgradeGoAhead>,
 	pub relay_dispatch_queue_size: Option<(u32, u32)>,
 	pub hrmp_ingress_channel_index: Option<Vec<ParaId>>,
 	pub hrmp_egress_channel_index: Option<Vec<ParaId>>,
@@ -59,6 +61,7 @@ impl Default for RelayStateSproofBuilder {
 				validation_upgrade_delay: 6,
 			},
 			dmq_mqc_head: None,
+			upgrade_go_ahead: None,
 			relay_dispatch_queue_size: None,
 			hrmp_ingress_channel_index: None,
 			hrmp_egress_channel_index: None,
@@ -68,6 +71,26 @@ impl Default for RelayStateSproofBuilder {
 	}
 }
 
+// TODO: derive `Copy` and `Clone` for `UpgradeGoAhead` to avoid manual implementation. 
+impl Clone for RelayStateSproofBuilder {
+	fn clone(&self) -> Self {
+		RelayStateSproofBuilder {
+			para_id: self.para_id,
+			host_config: self.host_config.clone(),
+			dmq_mqc_head: self.dmq_mqc_head.clone(),
+			upgrade_go_ahead: self.upgrade_go_ahead.as_ref().map(|u| match u {
+				UpgradeGoAhead::Abort => UpgradeGoAhead::Abort,
+				UpgradeGoAhead::GoAhead => UpgradeGoAhead::GoAhead,
+			}),
+			relay_dispatch_queue_size: self.relay_dispatch_queue_size,
+			hrmp_ingress_channel_index: self.hrmp_ingress_channel_index.clone(),
+			hrmp_egress_channel_index: self.hrmp_egress_channel_index.clone(),
+			hrmp_channels: self.hrmp_channels.clone(),
+			current_slot: self.current_slot.clone(),
+		}
+	}
+}
+
 impl RelayStateSproofBuilder {
 	/// Returns a mutable reference to HRMP channel metadata for a channel (`sender`, `self.para_id`).
 	///
@@ -129,6 +152,12 @@ impl RelayStateSproofBuilder {
 					relay_dispatch_queue_size.encode(),
 				);
 			}
+			if let Some(upgrade_go_ahead) = self.upgrade_go_ahead {
+				insert(
+					relay_chain::well_known_keys::upgrade_go_ahead_signal(self.para_id),
+					upgrade_go_ahead.encode(),
+				);
+			}
 			if let Some(hrmp_ingress_channel_index) = self.hrmp_ingress_channel_index {
 				let mut sorted = hrmp_ingress_channel_index.clone();
 				sorted.sort();