paras.rs

				})
				.for_each(sign_and_include_pvf_check_statement);

			// Check that `b` actually onboards.
			assert_eq!(<Paras as Store>::ActionsQueue::get(EXPECTED_SESSION + 2), vec![b]);

			// Check that the upgrade got scheduled.
			assert_eq!(
				<Paras as Store>::FutureCodeUpgrades::get(&a),
				Some(RELAY_PARENT + validation_upgrade_delay),
			);
		});
	}

	#[test]
	fn pvf_check_onboarding_reject_on_expiry() {
		let pvf_voting_ttl = 2;
		let a = ParaId::from(111);
		let validation_code: ValidationCode = vec![3, 2, 1].into();

		let genesis_config = MockGenesisConfig {
			configuration: crate::configuration::GenesisConfig {
				config: HostConfiguration {
					pvf_checking_enabled: true,
					pvf_voting_ttl,
					..Default::default()
				},
				..Default::default()
			},
			..Default::default()
		};

		new_test_ext(genesis_config).execute_with(|| {
			run_to_block(1, Some(vec![1]));

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

			// Make sure that we kicked off the PVF vote for this validation code and that the
			// validation code is stored.
			assert!(<Paras as Store>::PvfActiveVoteMap::get(&validation_code.hash()).is_some());
			check_code_is_stored(&validation_code);

			// Skip 2 sessions (i.e. `pvf_voting_ttl`) verifying that the code is still stored in
			// the intermediate session.
			assert_eq!(pvf_voting_ttl, 2);
			run_to_block(2, Some(vec![2]));
			check_code_is_stored(&validation_code);
			run_to_block(3, Some(vec![3]));

			// --- At this point the PVF vote for onboarding should be rejected.

			// Verify that the PVF is no longer stored and there is no active PVF vote.
			check_code_is_not_stored(&validation_code);
			assert!(<Paras as Store>::PvfActiveVoteMap::get(&validation_code.hash()).is_none());
			assert!(Paras::pvfs_require_precheck().is_empty());

			// Verify that at this point we can again try to initialize the same para.
			assert!(Paras::can_schedule_para_initialize(&a));
		});
	}

	#[test]
	fn pvf_check_upgrade_reject() {
		let a = ParaId::from(111);
		let new_code: ValidationCode = vec![3, 2, 1].into();

		let paras = vec![(
			a,
			ParaGenesisArgs {
				parachain: false,
				genesis_head: Default::default(),
				validation_code: ValidationCode(vec![]), // valid since in genesis
			},
		)];

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

		new_test_ext(genesis_config).execute_with(|| {
			// At this point `a` is already onboarded. Run to block 1 performing session change at
			// the end of block #0.
			run_to_block(2, Some(vec![1]));

			// Relay parent of the block that schedules the upgrade.
			const RELAY_PARENT: BlockNumber = 1;
			// Expected current session index.
			const EXPECTED_SESSION: SessionIndex = 1;

			Paras::schedule_code_upgrade(
				a,
				new_code.clone(),
				RELAY_PARENT,
				&Configuration::config(),
			);
			check_code_is_stored(&new_code);

			// Supermajority of validators vote against `new_code`. PVF should be rejected.
			IntoIterator::into_iter([0, 1, 2, 3])
				.map(|i| PvfCheckStatement {
					accept: false,
					subject: new_code.hash(),
					session_index: EXPECTED_SESSION,
					validator_index: i.into(),
				})
				.for_each(sign_and_include_pvf_check_statement);

			// Verify that the new code is discarded.
			check_code_is_not_stored(&new_code);

			assert!(<Paras as Store>::PvfActiveVoteMap::get(&new_code.hash()).is_none());
			assert!(Paras::pvfs_require_precheck().is_empty());
			assert!(<Paras as Store>::FutureCodeHash::get(&a).is_none());
		});
	}

	#[test]
	fn pvf_check_submit_vote() {
		let code_a: ValidationCode = vec![3, 2, 1].into();
		let code_b: ValidationCode = vec![1, 2, 3].into();

		let check = |stmt: PvfCheckStatement| -> (Result<_, _>, Result<_, _>) {
			let validators = &[
				Sr25519Keyring::Alice,
				Sr25519Keyring::Bob,
				Sr25519Keyring::Charlie,
				Sr25519Keyring::Dave,
				Sr25519Keyring::Ferdie,
				Sr25519Keyring::Eve, // <- this validator is not in the set
			];
			let signature: ValidatorSignature =
				validators[stmt.validator_index.0 as usize].sign(&stmt.signing_payload()).into();

			let call = Call::include_pvf_check_statement {
				stmt: stmt.clone(),
				signature: signature.clone(),
			};
			let validate_unsigned =
				<Paras as ValidateUnsigned>::validate_unsigned(TransactionSource::InBlock, &call)
					.map(|_| ());
			let dispatch_result =
				Paras::include_pvf_check_statement(None.into(), stmt.clone(), signature.clone());

			(validate_unsigned, dispatch_result)
		};

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

		new_test_ext(genesis_config).execute_with(|| {
			// Important to run this to seed the validators.
			run_to_block(1, Some(vec![1]));

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

			assert_eq!(
				check(PvfCheckStatement {
					accept: false,
					subject: code_a.hash(),
					session_index: 1,
					validator_index: 1.into(),
				}),
				(Ok(()), Ok(())),
			);

			// A vote in the same direction.
			let (unsigned, dispatch) = check(PvfCheckStatement {
				accept: false,
				subject: code_a.hash(),
				session_index: 1,
				validator_index: 1.into(),
			});
			assert_eq!(unsigned, Err(InvalidTransaction::Custom(INVALID_TX_DOUBLE_VOTE).into()));
			assert_err!(dispatch, Error::<Test>::PvfCheckDoubleVote);

			// Equivocation
			let (unsigned, dispatch) = check(PvfCheckStatement {
				accept: true,
				subject: code_a.hash(),
				session_index: 1,
				validator_index: 1.into(),
			});
			assert_eq!(unsigned, Err(InvalidTransaction::Custom(INVALID_TX_DOUBLE_VOTE).into()));
			assert_err!(dispatch, Error::<Test>::PvfCheckDoubleVote);

			// Vote for an earlier session.
			let (unsigned, dispatch) = check(PvfCheckStatement {
				accept: false,
				subject: code_a.hash(),
				session_index: 0,
				validator_index: 1.into(),
			});
			assert_eq!(unsigned, Err(InvalidTransaction::Stale.into()));
			assert_err!(dispatch, Error::<Test>::PvfCheckStatementStale);

			// Vote for an later session.
			let (unsigned, dispatch) = check(PvfCheckStatement {
				accept: false,
				subject: code_a.hash(),
				session_index: 2,
				validator_index: 1.into(),
			});
			assert_eq!(unsigned, Err(InvalidTransaction::Future.into()));
			assert_err!(dispatch, Error::<Test>::PvfCheckStatementFuture);

			// Validator not in the set.
			let (unsigned, dispatch) = check(PvfCheckStatement {
				accept: false,
				subject: code_a.hash(),
				session_index: 1,
				validator_index: 5.into(),
			});
			assert_eq!(
				unsigned,
				Err(InvalidTransaction::Custom(INVALID_TX_BAD_VALIDATOR_IDX).into())
			);
			assert_err!(dispatch, Error::<Test>::PvfCheckValidatorIndexOutOfBounds);

			// Bad subject (code_b)
			let (unsigned, dispatch) = check(PvfCheckStatement {
				accept: false,
				subject: code_b.hash(),
				session_index: 1,
				validator_index: 1.into(),
			});
			assert_eq!(unsigned, Err(InvalidTransaction::Custom(INVALID_TX_BAD_SUBJECT).into()));
			assert_err!(dispatch, Error::<Test>::PvfCheckSubjectInvalid);
		});
	}

	#[test]
	fn add_trusted_validation_code_inserts_with_no_users() {
		// This test is to ensure that trusted validation code is inserted into the storage
		// with the reference count equal to 0.
		let validation_code = ValidationCode(vec![1, 2, 3]);
		new_test_ext(Default::default()).execute_with(|| {
			assert_ok!(Paras::add_trusted_validation_code(Origin::root(), validation_code.clone()));
			assert_eq!(<Paras as Store>::CodeByHashRefs::get(&validation_code.hash()), 0,);
		});
	}

	#[test]
	fn add_trusted_validation_code_idempotent() {
		// This test makes sure that calling add_trusted_validation_code twice with the same
		// parameters is a no-op.
		let validation_code = ValidationCode(vec![1, 2, 3]);
		new_test_ext(Default::default()).execute_with(|| {
			assert_ok!(Paras::add_trusted_validation_code(Origin::root(), validation_code.clone()));
			assert_storage_noop!({
				assert_ok!(Paras::add_trusted_validation_code(
					Origin::root(),
					validation_code.clone()
				));
			});
		});
	}

	#[test]
	fn poke_unused_validation_code_removes_code_cleanly() {
		// This test makes sure that calling poke_unused_validation_code with a code that is currently
		// in the storage but has no users will remove it cleanly from the storage.
		let validation_code = ValidationCode(vec![1, 2, 3]);
		new_test_ext(Default::default()).execute_with(|| {
			assert_ok!(Paras::add_trusted_validation_code(Origin::root(), validation_code.clone()));
			assert_ok!(Paras::poke_unused_validation_code(Origin::root(), validation_code.hash()));

			assert_eq!(<Paras as Store>::CodeByHashRefs::get(&validation_code.hash()), 0);
			assert!(!<Paras as Store>::CodeByHash::contains_key(&validation_code.hash()));
		});
	}

	#[test]
	fn poke_unused_validation_code_doesnt_remove_code_with_users() {
		let para_id = 100.into();
		let validation_code = ValidationCode(vec![1, 2, 3]);
		new_test_ext(Default::default()).execute_with(|| {
			// First we add the code to the storage.
			assert_ok!(Paras::add_trusted_validation_code(Origin::root(), validation_code.clone()));

			// Then we add a user to the code, say by upgrading.
			run_to_block(2, None);
			Paras::schedule_code_upgrade(
				para_id,
				validation_code.clone(),
				1,
				&Configuration::config(),
			);
			Paras::note_new_head(para_id, HeadData::default(), 1);

			// Finally we poke the code, which should not remove it from the storage.
			assert_storage_noop!({
				assert_ok!(Paras::poke_unused_validation_code(
					Origin::root(),
					validation_code.hash()
				));
			});
			check_code_is_stored(&validation_code);
		});
	}

	#[test]
	fn increase_code_ref_doesnt_have_allergy_on_add_trusted_validation_code() {
		// Verify that accidential calling of increase_code_ref or decrease_code_ref does not lead
		// to a disaster.
		// NOTE that this test is extra paranoid, as it is not really possible to hit
		// `decrease_code_ref` without calling `increase_code_ref` first.
		let code = ValidationCode(vec![1, 2, 3]);

		new_test_ext(Default::default()).execute_with(|| {
			assert_ok!(Paras::add_trusted_validation_code(Origin::root(), code.clone()));
			Paras::increase_code_ref(&code.hash(), &code);
			Paras::increase_code_ref(&code.hash(), &code);
			assert!(<Paras as Store>::CodeByHash::contains_key(code.hash()));
			assert_eq!(<Paras as Store>::CodeByHashRefs::get(code.hash()), 2);
		});

		new_test_ext(Default::default()).execute_with(|| {
			assert_ok!(Paras::add_trusted_validation_code(Origin::root(), code.clone()));
			Paras::decrease_code_ref(&code.hash());
			assert!(<Paras as Store>::CodeByHash::contains_key(code.hash()));
			assert_eq!(<Paras as Store>::CodeByHashRefs::get(code.hash()), 0);
		});
	}

	#[test]
	fn add_trusted_validation_code_insta_approval() {
		// In particular, this tests that `kick_off_pvf_check` reacts to the `add_trusted_validation_code`
		// and uses the `CodeByHash::contains_key` which is what `add_trusted_validation_code` uses.
		let para_id = 100.into();
		let validation_code = ValidationCode(vec![1, 2, 3]);
		let validation_upgrade_delay = 25;
		let minimum_validation_upgrade_delay = 2;
		let genesis_config = MockGenesisConfig {
			configuration: crate::configuration::GenesisConfig {
				config: HostConfiguration {
					pvf_checking_enabled: true,
					validation_upgrade_delay,
					minimum_validation_upgrade_delay,
					..Default::default()
				},
				..Default::default()
			},
			..Default::default()
		};
		new_test_ext(genesis_config).execute_with(|| {
			assert_ok!(Paras::add_trusted_validation_code(Origin::root(), validation_code.clone()));

			// Then some parachain upgrades it's code with the relay-parent 1.
			run_to_block(2, None);
			Paras::schedule_code_upgrade(
				para_id,
				validation_code.clone(),
				1,
				&Configuration::config(),
			);
			Paras::note_new_head(para_id, HeadData::default(), 1);

			// Verify that the code upgrade has `expected_at` set to `26`. This is the behavior
			// equal to that of `pvf_checking_enabled: false`.
			assert_eq!(
				<Paras as Store>::FutureCodeUpgrades::get(&para_id),
				Some(1 + validation_upgrade_delay)
			);
		});
	}

	#[test]
	fn add_trusted_validation_code_enacts_existing_pvf_vote() {
		// This test makes sure that calling `add_trusted_validation_code` with a code that is
		// already going through PVF pre-checking voting will conclude the voting and enact the
		// code upgrade.
		let para_id = 100.into();
		let validation_code = ValidationCode(vec![1, 2, 3]);
		let validation_upgrade_delay = 25;
		let minimum_validation_upgrade_delay = 2;
		let genesis_config = MockGenesisConfig {
			configuration: crate::configuration::GenesisConfig {
				config: HostConfiguration {
					pvf_checking_enabled: true,
					validation_upgrade_delay,
					minimum_validation_upgrade_delay,
					..Default::default()
				},
				..Default::default()
			},
			..Default::default()
		};
		new_test_ext(genesis_config).execute_with(|| {
			// First, some parachain upgrades it's code with the relay-parent 1.
			run_to_block(2, None);
			Paras::schedule_code_upgrade(
				para_id,
				validation_code.clone(),
				1,
				&Configuration::config(),
			);
			Paras::note_new_head(para_id, HeadData::default(), 1);

			// No upgrade should be scheduled at this point. PVF pre-checking vote should run for
			// that PVF.
			assert!(<Paras as Store>::FutureCodeUpgrades::get(&para_id).is_none());
			assert!(<Paras as Store>::PvfActiveVoteMap::contains_key(&validation_code.hash()));

			// Then we add a trusted validation code. That should conclude the vote.
			assert_ok!(Paras::add_trusted_validation_code(Origin::root(), validation_code.clone()));
			assert!(<Paras as Store>::FutureCodeUpgrades::get(&para_id).is_some());
			assert!(!<Paras as Store>::PvfActiveVoteMap::contains_key(&validation_code.hash()));
		});
	}

	#[test]
	fn verify_upgrade_go_ahead_signal_is_externally_accessible() {
		use primitives::v1::well_known_keys;

		let a = ParaId::from(2020);

		new_test_ext(Default::default()).execute_with(|| {
			assert!(sp_io::storage::get(&well_known_keys::upgrade_go_ahead_signal(a)).is_none());
			<Paras as Store>::UpgradeGoAheadSignal::insert(&a, UpgradeGoAhead::GoAhead);
			assert_eq!(
				sp_io::storage::get(&well_known_keys::upgrade_go_ahead_signal(a)).unwrap(),
				vec![1u8],
			);
		});
	}

	#[test]
	fn verify_upgrade_restriction_signal_is_externally_accessible() {
		use primitives::v1::well_known_keys;

		let a = ParaId::from(2020);

		new_test_ext(Default::default()).execute_with(|| {
			assert!(sp_io::storage::get(&well_known_keys::upgrade_restriction_signal(a)).is_none());
			<Paras as Store>::UpgradeRestrictionSignal::insert(&a, UpgradeRestriction::Present);
			assert_eq!(
				sp_io::storage::get(&well_known_keys::upgrade_restriction_signal(a)).unwrap(),
				vec![0],
			);
		});
	}
}