paras.rs

			}
		})
	}

	/// Note that a para has progressed to a new head, where the new head was executed in the context
	/// of a relay-chain block with given number. This will apply pending code upgrades based
	/// on the relay-parent block number provided.
	pub(crate) fn note_new_head(
		id: ParaId,
		new_head: HeadData,
		execution_context: T::BlockNumber,
	) -> Weight {
		Heads::<T>::insert(&id, new_head);

		if let Some(expected_at) = <Self as Store>::FutureCodeUpgrades::get(&id) {
			if expected_at <= execution_context {
				<Self as Store>::FutureCodeUpgrades::remove(&id);
				<Self as Store>::UpgradeGoAheadSignal::remove(&id);

				// Both should always be `Some` in this case, since a code upgrade is scheduled.
				let new_code_hash = FutureCodeHash::<T>::take(&id).unwrap_or_default();
				let prior_code_hash = CurrentCodeHash::<T>::get(&id).unwrap_or_default();
				CurrentCodeHash::<T>::insert(&id, &new_code_hash);

				let log = ConsensusLog::ParaUpgradeCode(id, new_code_hash);
				<frame_system::Pallet<T>>::deposit_log(log.into());

				// `now` is only used for registering pruning as part of `fn note_past_code`
				let now = <frame_system::Pallet<T>>::block_number();

				let weight = Self::note_past_code(id, expected_at, now, prior_code_hash);

				// add 1 to writes due to heads update.
				weight + T::DbWeight::get().reads_writes(3, 1 + 3)
			} else {
				T::DbWeight::get().reads_writes(1, 1 + 0)
			}
		} else {
			T::DbWeight::get().reads_writes(1, 1)
		}
	}

	/// Fetches the validation code hash for the validation code to be used when validating a block
	/// in the context of the given relay-chain height. A second block number parameter may be used
	/// to tell the lookup to proceed as if an intermediate parablock has been with the given
	/// relay-chain height as its context. This may return the hash for the past, current, or
	/// (with certain choices of `assume_intermediate`) future code.
	///
	/// `assume_intermediate`, if provided, must be before `at`. This will return `None` if the validation
	/// code has been pruned.
	///
	/// To get associated code see [`Self::validation_code_at`].
	pub(crate) fn validation_code_hash_at(
		id: ParaId,
		at: T::BlockNumber,
		assume_intermediate: Option<T::BlockNumber>,
	) -> Option<ValidationCodeHash> {
		if assume_intermediate.as_ref().map_or(false, |i| &at <= i) {
			return None
		}

		let planned_upgrade = <Self as Store>::FutureCodeUpgrades::get(&id);
		let upgrade_applied_intermediate = match assume_intermediate {
			Some(a) => planned_upgrade.as_ref().map_or(false, |u| u <= &a),
			None => false,
		};

		if upgrade_applied_intermediate {
			FutureCodeHash::<T>::get(&id)
		} else {
			match Self::past_code_meta(&id).code_at(at) {
				None => None,
				Some(UseCodeAt::Current) => CurrentCodeHash::<T>::get(&id),
				Some(UseCodeAt::ReplacedAt(replaced)) =>
					<Self as Store>::PastCodeHash::get(&(id, replaced)),
			}
		}
	}

	/// Returns the current lifecycle state of the para.
	pub fn lifecycle(id: ParaId) -> Option<ParaLifecycle> {
		ParaLifecycles::<T>::get(&id)
	}

	/// Returns whether the given ID refers to a valid para.
	///
	/// Paras that are onboarding or offboarding are not included.
	pub fn is_valid_para(id: ParaId) -> bool {
		if let Some(state) = ParaLifecycles::<T>::get(&id) {
			!state.is_onboarding() && !state.is_offboarding()
		} else {
			false
		}
	}

	/// Whether a para ID corresponds to any live parachain.
	///
	/// Includes parachains which will downgrade to a parathread in the future.
	pub fn is_parachain(id: ParaId) -> bool {
		if let Some(state) = ParaLifecycles::<T>::get(&id) {
			state.is_parachain()
		} else {
			false
		}
	}

	/// Whether a para ID corresponds to any live parathread.
	///
	/// Includes parathreads which will upgrade to parachains in the future.
	pub fn is_parathread(id: ParaId) -> bool {
		if let Some(state) = ParaLifecycles::<T>::get(&id) {
			state.is_parathread()
		} else {
			false
		}
	}

	/// The block number of the last scheduled upgrade of the requested para. Includes future upgrades
	/// if the flag is set. This is the `expected_at` number, not the `activated_at` number.
	pub(crate) fn last_code_upgrade(id: ParaId, include_future: bool) -> Option<T::BlockNumber> {
		if include_future {
			if let Some(at) = Self::future_code_upgrade_at(id) {
				return Some(at)
			}
		}

		Self::past_code_meta(&id).most_recent_change()
	}

	/// Return the session index that should be used for any future scheduled changes.
	fn scheduled_session() -> SessionIndex {
		shared::Pallet::<T>::scheduled_session()
	}

	/// Store the validation code if not already stored, and increase the number of reference.
	///
	/// Returns the number of storage reads and number of storage writes.
	fn increase_code_ref(code_hash: &ValidationCodeHash, code: &ValidationCode) -> (u64, u64) {
		let reads = 1;
		let mut writes = 1;
		<Self as Store>::CodeByHashRefs::mutate(code_hash, |refs| {
			if *refs == 0 {
				writes += 1;
				<Self as Store>::CodeByHash::insert(code_hash, code);
			}
			*refs += 1;
		});
		(reads, writes)
	}

	/// Decrease the number of reference ofthe validation code and remove it from storage if zero
	/// is reached.
	fn decrease_code_ref(code_hash: &ValidationCodeHash) {
		let refs = <Self as Store>::CodeByHashRefs::get(code_hash);
		if refs <= 1 {
			<Self as Store>::CodeByHash::remove(code_hash);
			<Self as Store>::CodeByHashRefs::remove(code_hash);
		} else {
			<Self as Store>::CodeByHashRefs::insert(code_hash, refs - 1);
		}
	}

	/// Test function for triggering a new session in this pallet.
	#[cfg(any(feature = "std", feature = "runtime-benchmarks", test))]
	pub fn test_on_new_session() {
		Self::initializer_on_new_session(&SessionChangeNotification {
			session_index: shared::Pallet::<T>::session_index(),
			..Default::default()
		});
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use frame_support::assert_ok;
	use primitives::v1::BlockNumber;

	use crate::{
		configuration::HostConfiguration,
		mock::{new_test_ext, Configuration, MockGenesisConfig, Paras, ParasShared, System},
	};

	fn run_to_block(to: BlockNumber, new_session: Option<Vec<BlockNumber>>) {
		while System::block_number() < to {
			let b = System::block_number();
			Paras::initializer_finalize();
			ParasShared::initializer_finalize();
			if new_session.as_ref().map_or(false, |v| v.contains(&(b + 1))) {
				let mut session_change_notification = SessionChangeNotification::default();
				session_change_notification.session_index = ParasShared::session_index() + 1;
				ParasShared::initializer_on_new_session(
					session_change_notification.session_index,
					session_change_notification.random_seed,
					&session_change_notification.new_config,
					session_change_notification.validators.clone(),
				);
				Paras::initializer_on_new_session(&session_change_notification);
			}
			System::on_finalize(b);

			System::on_initialize(b + 1);
			System::set_block_number(b + 1);

			ParasShared::initializer_initialize(b + 1);
			Paras::initializer_initialize(b + 1);
		}
	}

	fn upgrade_at(
		expected_at: BlockNumber,
		activated_at: BlockNumber,
	) -> ReplacementTimes<BlockNumber> {
		ReplacementTimes { expected_at, activated_at }
	}

	fn check_code_is_stored(validation_code: &ValidationCode) {
		assert!(<Paras as Store>::CodeByHashRefs::get(validation_code.hash()) != 0);
		assert!(<Paras as Store>::CodeByHash::contains_key(validation_code.hash()));
	}

	fn check_code_is_not_stored(validation_code: &ValidationCode) {
		assert!(!<Paras as Store>::CodeByHashRefs::contains_key(validation_code.hash()));
		assert!(!<Paras as Store>::CodeByHash::contains_key(validation_code.hash()));
	}

	fn fetch_validation_code_at(
		para_id: ParaId,
		at: BlockNumber,
		assume_intermediate: Option<BlockNumber>,
	) -> Option<ValidationCode> {
		Paras::validation_code_hash_at(para_id, at, assume_intermediate)
			.and_then(Paras::code_by_hash)
	}

	#[test]
	fn para_past_code_meta_gives_right_code() {
		let mut past_code = ParaPastCodeMeta::default();
		assert_eq!(past_code.code_at(0u32), Some(UseCodeAt::Current));

		past_code.note_replacement(10, 12);
		assert_eq!(past_code.code_at(0), Some(UseCodeAt::ReplacedAt(10)));
		assert_eq!(past_code.code_at(10), Some(UseCodeAt::ReplacedAt(10)));
		assert_eq!(past_code.code_at(11), Some(UseCodeAt::ReplacedAt(10)));
		assert_eq!(past_code.code_at(12), Some(UseCodeAt::Current));

		past_code.note_replacement(20, 25);
		assert_eq!(past_code.code_at(1), Some(UseCodeAt::ReplacedAt(10)));
		assert_eq!(past_code.code_at(10), Some(UseCodeAt::ReplacedAt(10)));
		assert_eq!(past_code.code_at(11), Some(UseCodeAt::ReplacedAt(10)));
		assert_eq!(past_code.code_at(12), Some(UseCodeAt::ReplacedAt(20)));
		assert_eq!(past_code.code_at(24), Some(UseCodeAt::ReplacedAt(20)));
		assert_eq!(past_code.code_at(25), Some(UseCodeAt::Current));

		past_code.note_replacement(30, 30);
		assert_eq!(past_code.code_at(1), Some(UseCodeAt::ReplacedAt(10)));
		assert_eq!(past_code.code_at(10), Some(UseCodeAt::ReplacedAt(10)));
		assert_eq!(past_code.code_at(11), Some(UseCodeAt::ReplacedAt(10)));
		assert_eq!(past_code.code_at(12), Some(UseCodeAt::ReplacedAt(20)));
		assert_eq!(past_code.code_at(24), Some(UseCodeAt::ReplacedAt(20)));
		assert_eq!(past_code.code_at(25), Some(UseCodeAt::ReplacedAt(30)));
		assert_eq!(past_code.code_at(30), Some(UseCodeAt::Current));

		past_code.last_pruned = Some(5);
		assert_eq!(past_code.code_at(1), None);
		assert_eq!(past_code.code_at(5), None);
		assert_eq!(past_code.code_at(6), Some(UseCodeAt::ReplacedAt(10)));
		assert_eq!(past_code.code_at(24), Some(UseCodeAt::ReplacedAt(20)));
		assert_eq!(past_code.code_at(25), Some(UseCodeAt::ReplacedAt(30)));
		assert_eq!(past_code.code_at(30), Some(UseCodeAt::Current));
	}

	#[test]
	fn para_past_code_pruning_works_correctly() {
		let mut past_code = ParaPastCodeMeta::default();
		past_code.note_replacement(10u32, 10);
		past_code.note_replacement(20, 25);
		past_code.note_replacement(30, 35);

		let old = past_code.clone();
		assert!(past_code.prune_up_to(9).collect::<Vec<_>>().is_empty());
		assert_eq!(old, past_code);

		assert_eq!(past_code.prune_up_to(10).collect::<Vec<_>>(), vec![10]);
		assert_eq!(
			past_code,
			ParaPastCodeMeta {
				upgrade_times: vec![upgrade_at(20, 25), upgrade_at(30, 35)],
				last_pruned: Some(10),
			}
		);

		assert!(past_code.prune_up_to(21).collect::<Vec<_>>().is_empty());

		assert_eq!(past_code.prune_up_to(26).collect::<Vec<_>>(), vec![20]);
		assert_eq!(
			past_code,
			ParaPastCodeMeta { upgrade_times: vec![upgrade_at(30, 35)], last_pruned: Some(25) }
		);

		past_code.note_replacement(40, 42);
		past_code.note_replacement(50, 53);
		past_code.note_replacement(60, 66);

		assert_eq!(
			past_code,
			ParaPastCodeMeta {
				upgrade_times: vec![
					upgrade_at(30, 35),
					upgrade_at(40, 42),
					upgrade_at(50, 53),
					upgrade_at(60, 66)
				],
				last_pruned: Some(25),
			}
		);

		assert_eq!(past_code.prune_up_to(60).collect::<Vec<_>>(), vec![30, 40, 50]);
		assert_eq!(
			past_code,
			ParaPastCodeMeta { upgrade_times: vec![upgrade_at(60, 66)], last_pruned: Some(53) }
		);

		assert_eq!(past_code.most_recent_change(), Some(60));
		assert_eq!(past_code.prune_up_to(66).collect::<Vec<_>>(), vec![60]);

		assert_eq!(
			past_code,
			ParaPastCodeMeta { upgrade_times: Vec::new(), last_pruned: Some(66) }
		);
	}

	#[test]
	fn para_past_code_pruning_in_initialize() {
		let code_retention_period = 10;
		let paras = vec![
			(
				0u32.into(),
				ParaGenesisArgs {
					parachain: true,
					genesis_head: Default::default(),
					validation_code: Default::default(),
				},
			),
			(
				1u32.into(),
				ParaGenesisArgs {
					parachain: false,
					genesis_head: Default::default(),
					validation_code: Default::default(),
				},
			),
		];

		let genesis_config = MockGenesisConfig {
			paras: GenesisConfig { paras, ..Default::default() },
			configuration: crate::configuration::GenesisConfig {
				config: HostConfiguration { code_retention_period, ..Default::default() },
				..Default::default()
			},
			..Default::default()
		};

		new_test_ext(genesis_config).execute_with(|| {
			let id = ParaId::from(0u32);
			let at_block: BlockNumber = 10;
			let included_block: BlockNumber = 12;
			let validation_code = ValidationCode(vec![1, 2, 3]);

			Paras::increase_code_ref(&validation_code.hash(), &validation_code);
			<Paras as Store>::PastCodeHash::insert(&(id, at_block), &validation_code.hash());
			<Paras as Store>::PastCodePruning::put(&vec![(id, included_block)]);

			{
				let mut code_meta = Paras::past_code_meta(&id);
				code_meta.note_replacement(at_block, included_block);
				<Paras as Store>::PastCodeMeta::insert(&id, &code_meta);
			}

			let pruned_at: BlockNumber = included_block + code_retention_period + 1;
			assert_eq!(
				<Paras as Store>::PastCodeHash::get(&(id, at_block)),
				Some(validation_code.hash())
			);
			check_code_is_stored(&validation_code);

			run_to_block(pruned_at - 1, None);
			assert_eq!(
				<Paras as Store>::PastCodeHash::get(&(id, at_block)),
				Some(validation_code.hash())
			);
			assert_eq!(Paras::past_code_meta(&id).most_recent_change(), Some(at_block));
			check_code_is_stored(&validation_code);

			run_to_block(pruned_at, None);
			assert!(<Paras as Store>::PastCodeHash::get(&(id, at_block)).is_none());
			assert!(Paras::past_code_meta(&id).most_recent_change().is_none());
			check_code_is_not_stored(&validation_code);
		});
	}

	#[test]
	fn note_new_head_sets_head() {
		let code_retention_period = 10;
		let paras = vec![(
			0u32.into(),
			ParaGenesisArgs {
				parachain: true,
				genesis_head: Default::default(),
				validation_code: Default::default(),
			},
		)];

		let genesis_config = MockGenesisConfig {
			paras: GenesisConfig { paras, ..Default::default() },
			configuration: crate::configuration::GenesisConfig {
				config: HostConfiguration { code_retention_period, ..Default::default() },
				..Default::default()
			},
			..Default::default()
		};

		new_test_ext(genesis_config).execute_with(|| {
			let id_a = ParaId::from(0u32);

			assert_eq!(Paras::para_head(&id_a), Some(Default::default()));

			Paras::note_new_head(id_a, vec![1, 2, 3].into(), 0);

			assert_eq!(Paras::para_head(&id_a), Some(vec![1, 2, 3].into()));
		});
	}

	#[test]
	fn note_past_code_sets_up_pruning_correctly() {
		let code_retention_period = 10;
		let paras = vec![
			(
				0u32.into(),
				ParaGenesisArgs {
					parachain: true,
					genesis_head: Default::default(),
					validation_code: Default::default(),
				},
			),
			(
				1u32.into(),
				ParaGenesisArgs {
					parachain: false,
					genesis_head: Default::default(),
					validation_code: Default::default(),
				},
			),
		];

		let genesis_config = MockGenesisConfig {
			paras: GenesisConfig { paras, ..Default::default() },
			configuration: crate::configuration::GenesisConfig {
				config: HostConfiguration { code_retention_period, ..Default::default() },
				..Default::default()
			},
			..Default::default()
		};

		new_test_ext(genesis_config).execute_with(|| {
			let id_a = ParaId::from(0u32);
			let id_b = ParaId::from(1u32);

			Paras::note_past_code(id_a, 10, 12, ValidationCode(vec![1, 2, 3]).hash());
			Paras::note_past_code(id_b, 20, 23, ValidationCode(vec![4, 5, 6]).hash());

			assert_eq!(<Paras as Store>::PastCodePruning::get(), vec![(id_a, 12), (id_b, 23)]);
			assert_eq!(
				Paras::past_code_meta(&id_a),
				ParaPastCodeMeta { upgrade_times: vec![upgrade_at(10, 12)], last_pruned: None }
			);
			assert_eq!(
				Paras::past_code_meta(&id_b),
				ParaPastCodeMeta { upgrade_times: vec![upgrade_at(20, 23)], last_pruned: None }
			);
		});
	}

	#[test]
	fn code_upgrade_applied_after_delay() {
		let code_retention_period = 10;
		let validation_upgrade_delay = 5;
		let validation_upgrade_frequency = 10;

		let original_code = ValidationCode(vec![1, 2, 3]);
		let paras = vec![(
			0u32.into(),
			ParaGenesisArgs {
				parachain: true,
				genesis_head: Default::default(),
				validation_code: original_code.clone(),
			},
		)];

		let genesis_config = MockGenesisConfig {
			paras: GenesisConfig { paras, ..Default::default() },
			configuration: crate::configuration::GenesisConfig {
				config: HostConfiguration {
					code_retention_period,
					validation_upgrade_delay,
					validation_upgrade_frequency,
					..Default::default()
				},
				..Default::default()
			},
			..Default::default()
		};

		new_test_ext(genesis_config).execute_with(|| {
			check_code_is_stored(&original_code);

			let para_id = ParaId::from(0);
			let new_code = ValidationCode(vec![4, 5, 6]);

			run_to_block(2, None);
			assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));

			let expected_at = {
				// this parablock is in the context of block 1.
				let expected_at = 1 + validation_upgrade_delay;
				let next_possible_upgrade_at = 1 + validation_upgrade_frequency;
				Paras::schedule_code_upgrade(
					para_id,
					new_code.clone(),
					1,
					&Configuration::config(),
				);
				Paras::note_new_head(para_id, Default::default(), 1);

				assert!(Paras::past_code_meta(&para_id).most_recent_change().is_none());
				assert_eq!(<Paras as Store>::FutureCodeUpgrades::get(&para_id), Some(expected_at));
				assert_eq!(<Paras as Store>::FutureCodeHash::get(&para_id), Some(new_code.hash()));
				assert_eq!(<Paras as Store>::UpcomingUpgrades::get(), vec![(para_id, expected_at)]);
				assert_eq!(
					<Paras as Store>::UpgradeCooldowns::get(),
					vec![(para_id, next_possible_upgrade_at)]
				);
				assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));
				check_code_is_stored(&original_code);
				check_code_is_stored(&new_code);

				expected_at
			};

			run_to_block(expected_at, None);

			// the candidate is in the context of the parent of `expected_at`,
			// thus does not trigger the code upgrade.
			{
				Paras::note_new_head(para_id, Default::default(), expected_at - 1);

				assert!(Paras::past_code_meta(&para_id).most_recent_change().is_none());
				assert_eq!(<Paras as Store>::FutureCodeUpgrades::get(&para_id), Some(expected_at));
				assert_eq!(<Paras as Store>::FutureCodeHash::get(&para_id), Some(new_code.hash()));
				assert_eq!(
					<Paras as Store>::UpgradeGoAheadSignal::get(&para_id),
					Some(UpgradeGoAhead::GoAhead)
				);
				assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));
				check_code_is_stored(&original_code);
				check_code_is_stored(&new_code);
			}

			run_to_block(expected_at + 1, None);

			// the candidate is in the context of `expected_at`, and triggers
			// the upgrade.
			{
				Paras::note_new_head(para_id, Default::default(), expected_at);

				assert_eq!(Paras::past_code_meta(&para_id).most_recent_change(), Some(expected_at));
				assert_eq!(
					<Paras as Store>::PastCodeHash::get(&(para_id, expected_at)),
					Some(original_code.hash()),
				);
				assert!(<Paras as Store>::FutureCodeUpgrades::get(&para_id).is_none());
				assert!(<Paras as Store>::FutureCodeHash::get(&para_id).is_none());
				assert!(<Paras as Store>::UpgradeGoAheadSignal::get(&para_id).is_none());
				assert_eq!(Paras::current_code(&para_id), Some(new_code.clone()));
				check_code_is_stored(&original_code);
				check_code_is_stored(&new_code);
			}
		});
	}

	#[test]
	fn code_upgrade_applied_after_delay_even_when_late() {
		let code_retention_period = 10;
		let validation_upgrade_delay = 5;
		let validation_upgrade_frequency = 10;

		let original_code = ValidationCode(vec![1, 2, 3]);
		let paras = vec![(
			0u32.into(),
			ParaGenesisArgs {
				parachain: true,
				genesis_head: Default::default(),
				validation_code: original_code.clone(),
			},
		)];

		let genesis_config = MockGenesisConfig {
			paras: GenesisConfig { paras, ..Default::default() },
			configuration: crate::configuration::GenesisConfig {
				config: HostConfiguration {
					code_retention_period,
					validation_upgrade_delay,
					validation_upgrade_frequency,
					..Default::default()
				},
				..Default::default()
			},
			..Default::default()
		};

		new_test_ext(genesis_config).execute_with(|| {
			let para_id = ParaId::from(0);
			let new_code = ValidationCode(vec![4, 5, 6]);

			run_to_block(2, None);
			assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));

			let expected_at = {
				// this parablock is in the context of block 1.
				let expected_at = 1 + validation_upgrade_delay;
				let next_possible_upgrade_at = 1 + validation_upgrade_frequency;
				Paras::schedule_code_upgrade(
					para_id,
					new_code.clone(),
					1,
					&Configuration::config(),
				);
				Paras::note_new_head(para_id, Default::default(), 1);

				assert!(Paras::past_code_meta(&para_id).most_recent_change().is_none());
				assert_eq!(<Paras as Store>::FutureCodeUpgrades::get(&para_id), Some(expected_at));
				assert_eq!(<Paras as Store>::FutureCodeHash::get(&para_id), Some(new_code.hash()));
				assert_eq!(<Paras as Store>::UpcomingUpgrades::get(), vec![(para_id, expected_at)]);
				assert_eq!(
					<Paras as Store>::UpgradeCooldowns::get(),
					vec![(para_id, next_possible_upgrade_at)]
				);
				assert!(<Paras as Store>::UpgradeGoAheadSignal::get(&para_id).is_none());
				assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));

				expected_at
			};

			run_to_block(expected_at + 1 + 4, None);

			// the candidate is in the context of the first descendant of `expected_at`, and triggers
			// the upgrade.
			{
				// The signal should be set to go-ahead until the new head is actually processed.
				assert_eq!(
					<Paras as Store>::UpgradeGoAheadSignal::get(&para_id),
					Some(UpgradeGoAhead::GoAhead),
				);

				Paras::note_new_head(para_id, Default::default(), expected_at + 4);

				assert_eq!(Paras::past_code_meta(&para_id).most_recent_change(), Some(expected_at));

				// Some hypothetical block which would have triggered the code change
				// should still use the old code.
				assert_eq!(
					Paras::past_code_meta(&para_id).code_at(expected_at),
					Some(UseCodeAt::ReplacedAt(expected_at)),
				);

				// Some hypothetical block at the context which actually triggered the
				// code change should still use the old code.
				assert_eq!(
					Paras::past_code_meta(&para_id).code_at(expected_at + 4),
					Some(UseCodeAt::ReplacedAt(expected_at)),
				);

				// Some hypothetical block at the context after the code was upgraded
				// should use the new code.
				assert_eq!(
					Paras::past_code_meta(&para_id).code_at(expected_at + 4 + 1),
					Some(UseCodeAt::Current),
				);

				assert_eq!(
					<Paras as Store>::PastCodeHash::get(&(para_id, expected_at)),
					Some(original_code.hash()),
				);
				assert!(<Paras as Store>::FutureCodeUpgrades::get(&para_id).is_none());
				assert!(<Paras as Store>::FutureCodeHash::get(&para_id).is_none());
				assert!(<Paras as Store>::UpgradeGoAheadSignal::get(&para_id).is_none());
				assert_eq!(Paras::current_code(&para_id), Some(new_code.clone()));
			}
		});
	}

	#[test]
	fn submit_code_change_when_not_allowed_is_err() {
		let code_retention_period = 10;
		let validation_upgrade_delay = 7;

		let paras = vec![(
			0u32.into(),
			ParaGenesisArgs {
				parachain: true,
				genesis_head: Default::default(),
				validation_code: vec![1, 2, 3].into(),
			},
		)];

		let genesis_config = MockGenesisConfig {
			paras: GenesisConfig { paras, ..Default::default() },
			configuration: crate::configuration::GenesisConfig {
				config: HostConfiguration {
					code_retention_period,
					validation_upgrade_delay,
					..Default::default()
				},
				..Default::default()
			},
			..Default::default()
		};

		new_test_ext(genesis_config).execute_with(|| {
			let para_id = ParaId::from(0);
			let new_code = ValidationCode(vec![4, 5, 6]);
			let newer_code = ValidationCode(vec![4, 5, 6, 7]);

			run_to_block(1, None);
			Paras::schedule_code_upgrade(para_id, new_code.clone(), 1, &Configuration::config());
			assert_eq!(
				<Paras as Store>::FutureCodeUpgrades::get(&para_id),
				Some(1 + validation_upgrade_delay)
			);
			assert_eq!(<Paras as Store>::FutureCodeHash::get(&para_id), Some(new_code.hash()));
			check_code_is_stored(&new_code);

			// We expect that if an upgrade is signalled while there is already one pending we just
			// ignore it. Note that this is only true from perspective of this module.
			run_to_block(2, None);
			Paras::schedule_code_upgrade(para_id, newer_code.clone(), 2, &Configuration::config());
			assert_eq!(
				<Paras as Store>::FutureCodeUpgrades::get(&para_id),
				Some(1 + validation_upgrade_delay), // did not change since the same assertion from the last time.
			);
			assert_eq!(<Paras as Store>::FutureCodeHash::get(&para_id), Some(new_code.hash()));
			check_code_is_not_stored(&newer_code);
		});
	}

	#[test]
	fn full_parachain_cleanup_storage() {
		let code_retention_period = 10;
		let validation_upgrade_delay = 1 + 5;

		let original_code = ValidationCode(vec![1, 2, 3]);
		let paras = vec![(
			0u32.into(),
			ParaGenesisArgs {
				parachain: true,
				genesis_head: Default::default(),
				validation_code: original_code.clone(),
			},
		)];

		let genesis_config = MockGenesisConfig {
			paras: GenesisConfig { paras, ..Default::default() },
			configuration: crate::configuration::GenesisConfig {
				config: HostConfiguration {
					code_retention_period,
					validation_upgrade_delay,
					..Default::default()
				},
				..Default::default()
			},
			..Default::default()
		};

		new_test_ext(genesis_config).execute_with(|| {
			check_code_is_stored(&original_code);

			let para_id = ParaId::from(0);
			let new_code = ValidationCode(vec![4, 5, 6]);

			run_to_block(2, None);
			assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));
			check_code_is_stored(&original_code);

			let expected_at = {
				// this parablock is in the context of block 1.
				let expected_at = 1 + validation_upgrade_delay;
				Paras::schedule_code_upgrade(
					para_id,
					new_code.clone(),
					1,
					&Configuration::config(),
				);
				Paras::note_new_head(para_id, Default::default(), 1);

				assert!(Paras::past_code_meta(&para_id).most_recent_change().is_none());
				assert_eq!(<Paras as Store>::FutureCodeUpgrades::get(&para_id), Some(expected_at));
				assert_eq!(<Paras as Store>::FutureCodeHash::get(&para_id), Some(new_code.hash()));
				assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));
				check_code_is_stored(&original_code);
				check_code_is_stored(&new_code);

				expected_at
			};

			assert_ok!(Paras::schedule_para_cleanup(para_id));

			// Just scheduling cleanup shouldn't change anything.
			{
				assert_eq!(
					<Paras as Store>::ActionsQueue::get(Paras::scheduled_session()),
					vec![para_id],
				);
				assert_eq!(Paras::parachains(), vec![para_id]);

				assert!(Paras::past_code_meta(&para_id).most_recent_change().is_none());
				assert_eq!(<Paras as Store>::FutureCodeUpgrades::get(&para_id), Some(expected_at));
				assert_eq!(<Paras as Store>::FutureCodeHash::get(&para_id), Some(new_code.hash()));
				assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));
				check_code_is_stored(&original_code);
				check_code_is_stored(&new_code);

				assert_eq!(<Paras as Store>::Heads::get(&para_id), Some(Default::default()));
			}

			// run to block #4, with a 2 session changes at the end of the block 2 & 3.
			run_to_block(4, Some(vec![3, 4]));

			// cleaning up the parachain should place the current parachain code
			// into the past code buffer & schedule cleanup.
			assert_eq!(Paras::past_code_meta(&para_id).most_recent_change(), Some(3));
			assert_eq!(
				<Paras as Store>::PastCodeHash::get(&(para_id, 3)),
				Some(original_code.hash())
			);
			assert_eq!(<Paras as Store>::PastCodePruning::get(), vec![(para_id, 3)]);
			check_code_is_stored(&original_code);

			// any future upgrades haven't been used to validate yet, so those
			// are cleaned up immediately.
			assert!(<Paras as Store>::FutureCodeUpgrades::get(&para_id).is_none());
			assert!(<Paras as Store>::FutureCodeHash::get(&para_id).is_none());
			assert!(Paras::current_code(&para_id).is_none());
			check_code_is_not_stored(&new_code);

			// run to do the final cleanup
			let cleaned_up_at = 3 + code_retention_period + 1;
			run_to_block(cleaned_up_at, None);

			// now the final cleanup: last past code cleaned up, and this triggers meta cleanup.
			assert_eq!(Paras::past_code_meta(&para_id), Default::default());
			assert!(<Paras as Store>::PastCodeHash::get(&(para_id, 3)).is_none());
			assert!(<Paras as Store>::PastCodePruning::get().is_empty());
			check_code_is_not_stored(&original_code);
		});
	}

	#[test]
	fn para_incoming_at_session() {
		new_test_ext(Default::default()).execute_with(|| {
			run_to_block(1, None);

			let b = ParaId::from(525);
			let a = ParaId::from(999);
			let c = ParaId::from(333);

			assert_ok!(Paras::schedule_para_initialize(
				b,
				ParaGenesisArgs {
					parachain: true,
					genesis_head: vec![1].into(),
					validation_code: vec![1].into(),
				},
			));

			assert_ok!(Paras::schedule_para_initialize(
				a,
				ParaGenesisArgs {
					parachain: false,
					genesis_head: vec![2].into(),
					validation_code: vec![2].into(),
				},
			));

			assert_ok!(Paras::schedule_para_initialize(
				c,
				ParaGenesisArgs {
					parachain: true,
					genesis_head: vec![3].into(),
					validation_code: vec![3].into(),
				},
			));

			assert_eq!(
				<Paras as Store>::ActionsQueue::get(Paras::scheduled_session()),
				vec![c, b, a],
			);

			// Lifecycle is tracked correctly
			assert_eq!(<Paras as Store>::ParaLifecycles::get(&a), Some(ParaLifecycle::Onboarding));
			assert_eq!(<Paras as Store>::ParaLifecycles::get(&b), Some(ParaLifecycle::Onboarding));
			assert_eq!(<Paras as Store>::ParaLifecycles::get(&c), Some(ParaLifecycle::Onboarding));

			// run to block without session change.
			run_to_block(2, None);

			assert_eq!(Paras::parachains(), Vec::new());
			assert_eq!(
				<Paras as Store>::ActionsQueue::get(Paras::scheduled_session()),
				vec![c, b, a],
			);

			// Lifecycle is tracked correctly
			assert_eq!(<Paras as Store>::ParaLifecycles::get(&a), Some(ParaLifecycle::Onboarding));
			assert_eq!(<Paras as Store>::ParaLifecycles::get(&b), Some(ParaLifecycle::Onboarding));
			assert_eq!(<Paras as Store>::ParaLifecycles::get(&c), Some(ParaLifecycle::Onboarding));

			// Two sessions pass, so action queue is triggered
			run_to_block(4, Some(vec![3, 4]));

			assert_eq!(Paras::parachains(), vec![c, b]);
			assert_eq!(<Paras as Store>::ActionsQueue::get(Paras::scheduled_session()), Vec::new());

			// Lifecycle is tracked correctly
			assert_eq!(<Paras as Store>::ParaLifecycles::get(&a), Some(ParaLifecycle::Parathread));
			assert_eq!(<Paras as Store>::ParaLifecycles::get(&b), Some(ParaLifecycle::Parachain));
			assert_eq!(<Paras as Store>::ParaLifecycles::get(&c), Some(ParaLifecycle::Parachain));

			assert_eq!(Paras::current_code(&a), Some(vec![2].into()));
			assert_eq!(Paras::current_code(&b), Some(vec![1].into()));
			assert_eq!(Paras::current_code(&c), Some(vec![3].into()));
		})
	}

	#[test]
	fn code_hash_at_with_intermediate() {
		let code_retention_period = 10;
		let validation_upgrade_delay = 10;

		let paras = vec![(
			0u32.into(),
			ParaGenesisArgs {
				parachain: true,
				genesis_head: Default::default(),
				validation_code: vec![1, 2, 3].into(),
			},
		)];

		let genesis_config = MockGenesisConfig {
			paras: GenesisConfig { paras, ..Default::default() },
			configuration: crate::configuration::GenesisConfig {
				config: HostConfiguration {
					code_retention_period,
					validation_upgrade_delay,
					..Default::default()
				},
				..Default::default()
			},
			..Default::default()
		};

		new_test_ext(genesis_config).execute_with(|| {
			let para_id = ParaId::from(0);
			let old_code: ValidationCode = vec![1, 2, 3].into();
			let new_code: ValidationCode = vec![4, 5, 6].into();

			// expected_at = 10 = 0 + validation_upgrade_delay = 0 + 10
			Paras::schedule_code_upgrade(para_id, new_code.clone(), 0, &Configuration::config());
			assert_eq!(<Paras as Store>::FutureCodeUpgrades::get(&para_id), Some(10));

			// no intermediate, falls back on current/past.
			assert_eq!(fetch_validation_code_at(para_id, 1, None), Some(old_code.clone()));
			assert_eq!(fetch_validation_code_at(para_id, 10, None), Some(old_code.clone()));
			assert_eq!(fetch_validation_code_at(para_id, 100, None), Some(old_code.clone()));

			// intermediate before upgrade meant to be applied, falls back on current.
			assert_eq!(fetch_validation_code_at(para_id, 9, Some(8)), Some(old_code.clone()));
			assert_eq!(fetch_validation_code_at(para_id, 10, Some(9)), Some(old_code.clone()));
			assert_eq!(fetch_validation_code_at(para_id, 11, Some(9)), Some(old_code.clone()));

			// intermediate at or after upgrade applied
			assert_eq!(fetch_validation_code_at(para_id, 11, Some(10)), Some(new_code.clone()));
			assert_eq!(fetch_validation_code_at(para_id, 100, Some(11)), Some(new_code.clone()));

			run_to_block(code_retention_period + 5, None);

			// at <= intermediate not allowed
			assert_eq!(fetch_validation_code_at(para_id, 10, Some(10)), None);
			assert_eq!(fetch_validation_code_at(para_id, 9, Some(10)), None);
		});
	}