git.delta.rocks / unique-network / refs/commits / 8f7419bc1620

--- /dev/null
+++ b/pallets/collator-selection/CHANGELOG.md
@@ -0,0 +1,5 @@
+## [3.0.0] - 2022-10-28
+
+### Added
+
+ - Entire functionality moved over from cumulus/pallet-collator-selection (v3.0.0).
\ No newline at end of file

no changes

--- /dev/null
+++ b/pallets/collator-selection/src/benchmarking.rs
@@ -0,0 +1,275 @@
+// Copyright 2019-2022 Unique Network (Gibraltar) Ltd.
+// This file is part of Unique Network.
+
+// Unique Network is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Unique Network is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Unique Network. If not, see <http://www.gnu.org/licenses/>.
+
+// Original license:
+// Copyright (C) 2021 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Benchmarking setup for pallet-collator-selection
+
+use super::*;
+
+#[allow(unused)]
+use crate::Pallet as CollatorSelection;
+use frame_benchmarking::{account, benchmarks, impl_benchmark_test_suite, whitelisted_caller};
+use frame_support::{
+	assert_ok,
+	codec::Decode,
+	traits::{Currency, EnsureOrigin, Get},
+};
+use frame_system::{EventRecord, RawOrigin};
+use pallet_authorship::EventHandler;
+use pallet_session::{self as session, SessionManager};
+use sp_std::prelude::*;
+
+pub type BalanceOf<T> =
+	<<T as Config>::Currency as Currency<<T as frame_system::Config>::AccountId>>::Balance;
+
+const SEED: u32 = 0;
+
+// TODO: remove if this is given in substrate commit.
+macro_rules! whitelist {
+	($acc:ident) => {
+		frame_benchmarking::benchmarking::add_to_whitelist(
+			frame_system::Account::<T>::hashed_key_for(&$acc).into(),
+		);
+	};
+}
+
+fn assert_last_event<T: Config>(generic_event: <T as Config>::RuntimeEvent) {
+	let events = frame_system::Pallet::<T>::events();
+	let system_event: <T as frame_system::Config>::RuntimeEvent = generic_event.into();
+	// compare to the last event record
+	let EventRecord { event, .. } = &events[events.len() - 1];
+	assert_eq!(event, &system_event);
+}
+
+fn create_funded_user<T: Config>(
+	string: &'static str,
+	n: u32,
+	balance_factor: u32,
+) -> T::AccountId {
+	let user = account(string, n, SEED);
+	let balance = T::Currency::minimum_balance() * balance_factor.into();
+	let _ = T::Currency::make_free_balance_be(&user, balance);
+	user
+}
+
+fn keys<T: Config + session::Config>(c: u32) -> <T as session::Config>::Keys {
+	use rand::{RngCore, SeedableRng};
+
+	let keys = {
+		let mut keys = [0u8; 128];
+
+		if c > 0 {
+			let mut rng = rand::rngs::StdRng::seed_from_u64(c as u64);
+			rng.fill_bytes(&mut keys);
+		}
+
+		keys
+	};
+
+	Decode::decode(&mut &keys[..]).unwrap()
+}
+
+fn validator<T: Config + session::Config>(c: u32) -> (T::AccountId, <T as session::Config>::Keys) {
+	(create_funded_user::<T>("candidate", c, 1000), keys::<T>(c))
+}
+
+fn register_validators<T: Config + session::Config>(count: u32) -> Vec<T::AccountId> {
+	let validators = (0..count).map(|c| validator::<T>(c)).collect::<Vec<_>>();
+
+	for (who, keys) in validators.clone() {
+		<session::Pallet<T>>::set_keys(RawOrigin::Signed(who).into(), keys, Vec::new()).unwrap();
+	}
+
+	validators.into_iter().map(|(who, _)| who).collect()
+}
+
+fn register_candidates<T: Config>(count: u32) {
+	let candidates = (0..count).map(|c| account("candidate", c, SEED)).collect::<Vec<_>>();
+	assert!(<CandidacyBond<T>>::get() > 0u32.into(), "Bond cannot be zero!");
+
+	for who in candidates {
+		T::Currency::make_free_balance_be(&who, <CandidacyBond<T>>::get() * 2u32.into());
+		<CollatorSelection<T>>::register_as_candidate(RawOrigin::Signed(who).into()).unwrap();
+	}
+}
+
+benchmarks! {
+	where_clause { where T: pallet_authorship::Config + session::Config }
+
+	set_invulnerables {
+		let b in 1 .. T::MaxInvulnerables::get();
+		let new_invulnerables = register_validators::<T>(b);
+		let origin = T::UpdateOrigin::successful_origin();
+	}: {
+		assert_ok!(
+			<CollatorSelection<T>>::set_invulnerables(origin, new_invulnerables.clone())
+		);
+	}
+	verify {
+		assert_last_event::<T>(Event::NewInvulnerables{invulnerables: new_invulnerables}.into());
+	}
+
+	set_desired_candidates {
+		let max: u32 = 999;
+		let origin = T::UpdateOrigin::successful_origin();
+	}: {
+		assert_ok!(
+			<CollatorSelection<T>>::set_desired_candidates(origin, max.clone())
+		);
+	}
+	verify {
+		assert_last_event::<T>(Event::NewDesiredCandidates{desired_candidates: max}.into());
+	}
+
+	set_candidacy_bond {
+		let bond_amount: BalanceOf<T> = T::Currency::minimum_balance() * 10u32.into();
+		let origin = T::UpdateOrigin::successful_origin();
+	}: {
+		assert_ok!(
+			<CollatorSelection<T>>::set_candidacy_bond(origin, bond_amount.clone())
+		);
+	}
+	verify {
+		assert_last_event::<T>(Event::NewCandidacyBond{bond_amount}.into());
+	}
+
+	// worse case is when we have all the max-candidate slots filled except one, and we fill that
+	// one.
+	register_as_candidate {
+		let c in 1 .. T::MaxCandidates::get();
+
+		<CandidacyBond<T>>::put(T::Currency::minimum_balance());
+		<DesiredCandidates<T>>::put(c + 1);
+
+		register_validators::<T>(c);
+		register_candidates::<T>(c);
+
+		let caller: T::AccountId = whitelisted_caller();
+		let bond: BalanceOf<T> = T::Currency::minimum_balance() * 2u32.into();
+		T::Currency::make_free_balance_be(&caller, bond.clone());
+
+		<session::Pallet<T>>::set_keys(
+			RawOrigin::Signed(caller.clone()).into(),
+			keys::<T>(c + 1),
+			Vec::new()
+		).unwrap();
+
+	}: _(RawOrigin::Signed(caller.clone()))
+	verify {
+		assert_last_event::<T>(Event::CandidateAdded{account_id: caller, deposit: bond / 2u32.into()}.into());
+	}
+
+	// worse case is the last candidate leaving.
+	leave_intent {
+		let c in (T::MinCandidates::get() + 1) .. T::MaxCandidates::get();
+		<CandidacyBond<T>>::put(T::Currency::minimum_balance());
+		<DesiredCandidates<T>>::put(c);
+
+		register_validators::<T>(c);
+		register_candidates::<T>(c);
+
+		let leaving = <Candidates<T>>::get().last().unwrap().who.clone();
+		whitelist!(leaving);
+	}: _(RawOrigin::Signed(leaving.clone()))
+	verify {
+		assert_last_event::<T>(Event::CandidateRemoved{account_id: leaving}.into());
+	}
+
+	// worse case is paying a non-existing candidate account.
+	note_author {
+		<CandidacyBond<T>>::put(T::Currency::minimum_balance());
+		T::Currency::make_free_balance_be(
+			&<CollatorSelection<T>>::account_id(),
+			T::Currency::minimum_balance() * 4u32.into(),
+		);
+		let author = account("author", 0, SEED);
+		let new_block: T::BlockNumber = 10u32.into();
+
+		frame_system::Pallet::<T>::set_block_number(new_block);
+		assert!(T::Currency::free_balance(&author) == 0u32.into());
+	}: {
+		<CollatorSelection<T> as EventHandler<_, _>>::note_author(author.clone())
+	} verify {
+		assert!(T::Currency::free_balance(&author) > 0u32.into());
+		assert_eq!(frame_system::Pallet::<T>::block_number(), new_block);
+	}
+
+	// worst case for new session.
+	new_session {
+		let r in 1 .. T::MaxCandidates::get();
+		let c in 1 .. T::MaxCandidates::get();
+
+		<CandidacyBond<T>>::put(T::Currency::minimum_balance());
+		<DesiredCandidates<T>>::put(c);
+		frame_system::Pallet::<T>::set_block_number(0u32.into());
+
+		register_validators::<T>(c);
+		register_candidates::<T>(c);
+
+		let new_block: T::BlockNumber = 1800u32.into();
+		let zero_block: T::BlockNumber = 0u32.into();
+		let candidates = <Candidates<T>>::get();
+
+		let non_removals = c.saturating_sub(r);
+
+		for i in 0..c {
+			<LastAuthoredBlock<T>>::insert(candidates[i as usize].who.clone(), zero_block);
+		}
+
+		if non_removals > 0 {
+			for i in 0..non_removals {
+				<LastAuthoredBlock<T>>::insert(candidates[i as usize].who.clone(), new_block);
+			}
+		} else {
+			for i in 0..c {
+				<LastAuthoredBlock<T>>::insert(candidates[i as usize].who.clone(), new_block);
+			}
+		}
+
+		let pre_length = <Candidates<T>>::get().len();
+
+		frame_system::Pallet::<T>::set_block_number(new_block);
+
+		assert!(<Candidates<T>>::get().len() == c as usize);
+	}: {
+		<CollatorSelection<T> as SessionManager<_>>::new_session(0)
+	} verify {
+		if c > r && non_removals >= T::MinCandidates::get() {
+			assert!(<Candidates<T>>::get().len() < pre_length);
+		} else if c > r && non_removals < T::MinCandidates::get() {
+			assert!(<Candidates<T>>::get().len() == T::MinCandidates::get() as usize);
+		} else {
+			assert!(<Candidates<T>>::get().len() == pre_length);
+		}
+	}
+}
+
+impl_benchmark_test_suite!(CollatorSelection, crate::mock::new_test_ext(), crate::mock::Test,);

--- /dev/null
+++ b/pallets/collator-selection/src/lib.rs
@@ -0,0 +1,545 @@
+// Copyright 2019-2022 Unique Network (Gibraltar) Ltd.
+// This file is part of Unique Network.
+
+// Unique Network is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Unique Network is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Unique Network. If not, see <http://www.gnu.org/licenses/>.
+
+// Original license:
+// Copyright (C) 2021 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Collator Selection pallet.
+//!
+//! A pallet to manage collators in a parachain.
+//!
+//! ## Overview
+//!
+//! The Collator Selection pallet manages the collators of a parachain. **Collation is _not_ a
+//! secure activity** and this pallet does not implement any game-theoretic mechanisms to meet BFT
+//! safety assumptions of the chosen set.
+//!
+//! ## Terminology
+//!
+//! - Collator: A parachain block producer.
+//! - Bond: An amount of `Balance` _reserved_ for candidate registration.
+//! - Invulnerable: An account guaranteed to be in the collator set.
+//!
+//! ## Implementation
+//!
+//! The final `Collators` are aggregated from two individual lists:
+//!
+//! 1. [`Invulnerables`]: a set of collators appointed by governance. These accounts will always be
+//!    collators.
+//! 2. [`Candidates`]: these are *candidates to the collation task* and may or may not be elected as
+//!    a final collator.
+//!
+//! The current implementation resolves congestion of [`Candidates`] in a first-come-first-serve
+//! manner.
+//!
+//! Candidates will not be allowed to get kicked or leave_intent if the total number of candidates
+//! fall below MinCandidates. This is for potential disaster recovery scenarios.
+//!
+//! ### Rewards
+//!
+//! The Collator Selection pallet maintains an on-chain account (the "Pot"). In each block, the
+//! collator who authored it receives:
+//!
+//! - Half the value of the Pot.
+//! - Half the value of the transaction fees within the block. The other half of the transaction
+//!   fees are deposited into the Pot.
+//!
+//! To initiate rewards an ED needs to be transferred to the pot address.
+//!
+//! Note: Eventually the Pot distribution may be modified as discussed in
+//! [this issue](https://github.com/paritytech/statemint/issues/21#issuecomment-810481073).
+
+#![cfg_attr(not(feature = "std"), no_std)]
+
+pub use pallet::*;
+
+#[cfg(test)]
+mod mock;
+
+#[cfg(test)]
+mod tests;
+
+#[cfg(feature = "runtime-benchmarks")]
+mod benchmarking;
+pub mod weights;
+
+#[frame_support::pallet]
+pub mod pallet {
+	pub use crate::weights::WeightInfo;
+	use core::ops::Div;
+	use frame_support::{
+		dispatch::{DispatchClass, DispatchResultWithPostInfo},
+		inherent::Vec,
+		pallet_prelude::*,
+		sp_runtime::{
+			traits::{AccountIdConversion, CheckedSub, Saturating, Zero},
+			RuntimeDebug,
+		},
+		traits::{
+			Currency, EnsureOrigin, ExistenceRequirement::KeepAlive, ReservableCurrency,
+			ValidatorRegistration,
+		},
+		BoundedVec, PalletId,
+	};
+	use frame_system::{pallet_prelude::*, Config as SystemConfig};
+	use pallet_session::SessionManager;
+	use sp_runtime::traits::Convert;
+	use sp_staking::SessionIndex;
+
+	type BalanceOf<T> =
+		<<T as Config>::Currency as Currency<<T as SystemConfig>::AccountId>>::Balance;
+
+	/// A convertor from collators id. Since this pallet does not have stash/controller, this is
+	/// just identity.
+	pub struct IdentityCollator;
+	impl<T> sp_runtime::traits::Convert<T, Option<T>> for IdentityCollator {
+		fn convert(t: T) -> Option<T> {
+			Some(t)
+		}
+	}
+
+	/// Configure the pallet by specifying the parameters and types on which it depends.
+	#[pallet::config]
+	pub trait Config: frame_system::Config {
+		/// Overarching event type.
+		type RuntimeEvent: From<Event<Self>> + IsType<<Self as frame_system::Config>::RuntimeEvent>;
+
+		/// The currency mechanism.
+		type Currency: ReservableCurrency<Self::AccountId>;
+
+		/// Origin that can dictate updating parameters of this pallet.
+		type UpdateOrigin: EnsureOrigin<Self::RuntimeOrigin>;
+
+		/// Account Identifier from which the internal Pot is generated.
+		type PotId: Get<PalletId>;
+
+		/// Maximum number of candidates that we should have. This is enforced in code.
+		///
+		/// This does not take into account the invulnerables.
+		type MaxCandidates: Get<u32>;
+
+		/// Minimum number of candidates that we should have. This is used for disaster recovery.
+		///
+		/// This does not take into account the invulnerables.
+		type MinCandidates: Get<u32>;
+
+		/// Maximum number of invulnerables. This is enforced in code.
+		type MaxInvulnerables: Get<u32>;
+
+		// Will be kicked if block is not produced in threshold.
+		type KickThreshold: Get<Self::BlockNumber>;
+
+		/// A stable ID for a validator.
+		type ValidatorId: Member + Parameter;
+
+		/// A conversion from account ID to validator ID.
+		///
+		/// Its cost must be at most one storage read.
+		type ValidatorIdOf: Convert<Self::AccountId, Option<Self::ValidatorId>>;
+
+		/// Validate a user is registered
+		type ValidatorRegistration: ValidatorRegistration<Self::ValidatorId>;
+
+		/// The weight information of this pallet.
+		type WeightInfo: WeightInfo;
+	}
+
+	/// Basic information about a collation candidate.
+	#[derive(
+		PartialEq, Eq, Clone, Encode, Decode, RuntimeDebug, scale_info::TypeInfo, MaxEncodedLen,
+	)]
+	pub struct CandidateInfo<AccountId, Balance> {
+		/// Account identifier.
+		pub who: AccountId,
+		/// Reserved deposit.
+		pub deposit: Balance,
+	}
+
+	#[pallet::pallet]
+	#[pallet::generate_store(pub(super) trait Store)]
+	pub struct Pallet<T>(_);
+
+	/// The invulnerable, fixed collators.
+	#[pallet::storage]
+	#[pallet::getter(fn invulnerables)]
+	pub type Invulnerables<T: Config> =
+		StorageValue<_, BoundedVec<T::AccountId, T::MaxInvulnerables>, ValueQuery>;
+
+	/// The (community, limited) collation candidates.
+	#[pallet::storage]
+	#[pallet::getter(fn candidates)]
+	pub type Candidates<T: Config> = StorageValue<
+		_,
+		BoundedVec<CandidateInfo<T::AccountId, BalanceOf<T>>, T::MaxCandidates>,
+		ValueQuery,
+	>;
+
+	/// Last block authored by collator.
+	#[pallet::storage]
+	#[pallet::getter(fn last_authored_block)]
+	pub type LastAuthoredBlock<T: Config> =
+		StorageMap<_, Twox64Concat, T::AccountId, T::BlockNumber, ValueQuery>;
+
+	/// Desired number of candidates.
+	///
+	/// This should ideally always be less than [`Config::MaxCandidates`] for weights to be correct.
+	#[pallet::storage]
+	#[pallet::getter(fn desired_candidates)]
+	pub type DesiredCandidates<T> = StorageValue<_, u32, ValueQuery>;
+
+	/// Fixed amount to deposit to become a collator.
+	///
+	/// When a collator calls `leave_intent` they immediately receive the deposit back.
+	#[pallet::storage]
+	#[pallet::getter(fn candidacy_bond)]
+	pub type CandidacyBond<T> = StorageValue<_, BalanceOf<T>, ValueQuery>;
+
+	#[pallet::genesis_config]
+	pub struct GenesisConfig<T: Config> {
+		pub invulnerables: Vec<T::AccountId>,
+		pub candidacy_bond: BalanceOf<T>,
+		pub desired_candidates: u32,
+	}
+
+	#[cfg(feature = "std")]
+	impl<T: Config> Default for GenesisConfig<T> {
+		fn default() -> Self {
+			Self {
+				invulnerables: Default::default(),
+				candidacy_bond: Default::default(),
+				desired_candidates: Default::default(),
+			}
+		}
+	}
+
+	#[pallet::genesis_build]
+	impl<T: Config> GenesisBuild<T> for GenesisConfig<T> {
+		fn build(&self) {
+			let duplicate_invulnerables =
+				self.invulnerables.iter().collect::<std::collections::BTreeSet<_>>();
+			assert!(
+				duplicate_invulnerables.len() == self.invulnerables.len(),
+				"duplicate invulnerables in genesis."
+			);
+
+			let bounded_invulnerables =
+				BoundedVec::<_, T::MaxInvulnerables>::try_from(self.invulnerables.clone())
+					.expect("genesis invulnerables are more than T::MaxInvulnerables");
+			assert!(
+				T::MaxCandidates::get() >= self.desired_candidates,
+				"genesis desired_candidates are more than T::MaxCandidates",
+			);
+
+			<DesiredCandidates<T>>::put(&self.desired_candidates);
+			<CandidacyBond<T>>::put(&self.candidacy_bond);
+			<Invulnerables<T>>::put(bounded_invulnerables);
+		}
+	}
+
+	#[pallet::event]
+	#[pallet::generate_deposit(pub(super) fn deposit_event)]
+	pub enum Event<T: Config> {
+		NewInvulnerables { invulnerables: Vec<T::AccountId> },
+		NewDesiredCandidates { desired_candidates: u32 },
+		NewCandidacyBond { bond_amount: BalanceOf<T> },
+		CandidateAdded { account_id: T::AccountId, deposit: BalanceOf<T> },
+		CandidateRemoved { account_id: T::AccountId },
+	}
+
+	// Errors inform users that something went wrong.
+	#[pallet::error]
+	pub enum Error<T> {
+		/// Too many candidates
+		TooManyCandidates,
+		/// Too few candidates
+		TooFewCandidates,
+		/// Unknown error
+		Unknown,
+		/// Permission issue
+		Permission,
+		/// User is already a candidate
+		AlreadyCandidate,
+		/// User is not a candidate
+		NotCandidate,
+		/// Too many invulnerables
+		TooManyInvulnerables,
+		/// User is already an Invulnerable
+		AlreadyInvulnerable,
+		/// Account has no associated validator ID
+		NoAssociatedValidatorId,
+		/// Validator ID is not yet registered
+		ValidatorNotRegistered,
+	}
+
+	#[pallet::hooks]
+	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {}
+
+	#[pallet::call]
+	impl<T: Config> Pallet<T> {
+		/// Set the list of invulnerable (fixed) collators.
+		#[pallet::weight(T::WeightInfo::set_invulnerables(new.len() as u32))]
+		pub fn set_invulnerables(
+			origin: OriginFor<T>,
+			new: Vec<T::AccountId>,
+		) -> DispatchResultWithPostInfo {
+			T::UpdateOrigin::ensure_origin(origin)?;
+			let bounded_invulnerables = BoundedVec::<_, T::MaxInvulnerables>::try_from(new)
+				.map_err(|_| Error::<T>::TooManyInvulnerables)?;
+
+			// check if the invulnerables have associated validator keys before they are set
+			for account_id in bounded_invulnerables.iter() {
+				let validator_key = T::ValidatorIdOf::convert(account_id.clone())
+					.ok_or(Error::<T>::NoAssociatedValidatorId)?;
+				ensure!(
+					T::ValidatorRegistration::is_registered(&validator_key),
+					Error::<T>::ValidatorNotRegistered
+				);
+			}
+
+			<Invulnerables<T>>::put(&bounded_invulnerables);
+			Self::deposit_event(Event::NewInvulnerables {
+				invulnerables: bounded_invulnerables.to_vec(),
+			});
+			Ok(().into())
+		}
+
+		/// Set the ideal number of collators (not including the invulnerables).
+		/// If lowering this number, then the number of running collators could be higher than this figure.
+		/// Aside from that edge case, there should be no other way to have more collators than the desired number.
+		#[pallet::weight(T::WeightInfo::set_desired_candidates())]
+		pub fn set_desired_candidates(
+			origin: OriginFor<T>,
+			max: u32,
+		) -> DispatchResultWithPostInfo {
+			T::UpdateOrigin::ensure_origin(origin)?;
+			// we trust origin calls, this is just a for more accurate benchmarking
+			if max > T::MaxCandidates::get() {
+				log::warn!("max > T::MaxCandidates; you might need to run benchmarks again");
+			}
+			<DesiredCandidates<T>>::put(&max);
+			Self::deposit_event(Event::NewDesiredCandidates { desired_candidates: max });
+			Ok(().into())
+		}
+
+		/// Set the candidacy bond amount.
+		#[pallet::weight(T::WeightInfo::set_candidacy_bond())]
+		pub fn set_candidacy_bond(
+			origin: OriginFor<T>,
+			bond: BalanceOf<T>,
+		) -> DispatchResultWithPostInfo {
+			T::UpdateOrigin::ensure_origin(origin)?;
+			<CandidacyBond<T>>::put(&bond);
+			Self::deposit_event(Event::NewCandidacyBond { bond_amount: bond });
+			Ok(().into())
+		}
+
+		/// Register this account as a collator candidate. The account must (a) already have
+		/// registered session keys and (b) be able to reserve the `CandidacyBond`.
+		///
+		/// This call is not available to `Invulnerable` collators.
+		#[pallet::weight(T::WeightInfo::register_as_candidate(T::MaxCandidates::get()))]
+		pub fn register_as_candidate(origin: OriginFor<T>) -> DispatchResultWithPostInfo {
+			let who = ensure_signed(origin)?;
+
+			// ensure we are below limit.
+			let length = <Candidates<T>>::decode_len().unwrap_or_default();
+			ensure!((length as u32) < Self::desired_candidates(), Error::<T>::TooManyCandidates);
+			ensure!(!Self::invulnerables().contains(&who), Error::<T>::AlreadyInvulnerable);
+
+			let validator_key = T::ValidatorIdOf::convert(who.clone())
+				.ok_or(Error::<T>::NoAssociatedValidatorId)?;
+			ensure!(
+				T::ValidatorRegistration::is_registered(&validator_key),
+				Error::<T>::ValidatorNotRegistered
+			);
+
+			let deposit = Self::candidacy_bond();
+			// First authored block is current block plus kick threshold to handle session delay
+			let incoming = CandidateInfo { who: who.clone(), deposit };
+
+			let current_count =
+				<Candidates<T>>::try_mutate(|candidates| -> Result<usize, DispatchError> {
+					if candidates.iter().any(|candidate| candidate.who == who) {
+						Err(Error::<T>::AlreadyCandidate)?
+					} else {
+						T::Currency::reserve(&who, deposit)?;
+						candidates.try_push(incoming).map_err(|_| Error::<T>::TooManyCandidates)?;
+						<LastAuthoredBlock<T>>::insert(
+							who.clone(),
+							frame_system::Pallet::<T>::block_number() + T::KickThreshold::get(),
+						);
+						Ok(candidates.len())
+					}
+				})?;
+
+			Self::deposit_event(Event::CandidateAdded { account_id: who, deposit });
+			Ok(Some(T::WeightInfo::register_as_candidate(current_count as u32)).into())
+		}
+
+		/// Deregister `origin` as a collator candidate. Note that the collator can only leave on
+		/// session change. The `CandidacyBond` will be unreserved immediately.
+		///
+		/// This call will fail if the total number of candidates would drop below `MinCandidates`.
+		///
+		/// This call is not available to `Invulnerable` collators.
+		#[pallet::weight(T::WeightInfo::leave_intent(T::MaxCandidates::get()))]
+		pub fn leave_intent(origin: OriginFor<T>) -> DispatchResultWithPostInfo {
+			let who = ensure_signed(origin)?;
+			ensure!(
+				Self::candidates().len() as u32 > T::MinCandidates::get(),
+				Error::<T>::TooFewCandidates
+			);
+			let current_count = Self::try_remove_candidate(&who)?;
+
+			Ok(Some(T::WeightInfo::leave_intent(current_count as u32)).into())
+		}
+	}
+
+	impl<T: Config> Pallet<T> {
+		/// Get a unique, inaccessible account id from the `PotId`.
+		pub fn account_id() -> T::AccountId {
+			T::PotId::get().into_account_truncating()
+		}
+
+		/// Removes a candidate if they exist and sends them back their deposit
+		fn try_remove_candidate(who: &T::AccountId) -> Result<usize, DispatchError> {
+			let current_count =
+				<Candidates<T>>::try_mutate(|candidates| -> Result<usize, DispatchError> {
+					let index = candidates
+						.iter()
+						.position(|candidate| candidate.who == *who)
+						.ok_or(Error::<T>::NotCandidate)?;
+					let candidate = candidates.remove(index);
+					T::Currency::unreserve(who, candidate.deposit);
+					<LastAuthoredBlock<T>>::remove(who.clone());
+					Ok(candidates.len())
+				})?;
+			Self::deposit_event(Event::CandidateRemoved { account_id: who.clone() });
+			Ok(current_count)
+		}
+
+		/// Assemble the current set of candidates and invulnerables into the next collator set.
+		///
+		/// This is done on the fly, as frequent as we are told to do so, as the session manager.
+		pub fn assemble_collators(
+			candidates: BoundedVec<T::AccountId, T::MaxCandidates>,
+		) -> Vec<T::AccountId> {
+			let mut collators = Self::invulnerables().to_vec();
+			collators.extend(candidates);
+			collators
+		}
+
+		/// Kicks out candidates that did not produce a block in the kick threshold
+		/// and refund their deposits.
+		pub fn kick_stale_candidates(
+			candidates: BoundedVec<CandidateInfo<T::AccountId, BalanceOf<T>>, T::MaxCandidates>,
+		) -> BoundedVec<T::AccountId, T::MaxCandidates> {
+			let now = frame_system::Pallet::<T>::block_number();
+			let kick_threshold = T::KickThreshold::get();
+			candidates
+				.into_iter()
+				.filter_map(|c| {
+					let last_block = <LastAuthoredBlock<T>>::get(c.who.clone());
+					let since_last = now.saturating_sub(last_block);
+					if since_last < kick_threshold ||
+						Self::candidates().len() as u32 <= T::MinCandidates::get()
+					{
+						Some(c.who)
+					} else {
+						let outcome = Self::try_remove_candidate(&c.who);
+						if let Err(why) = outcome {
+							log::warn!("Failed to remove candidate {:?}", why);
+							debug_assert!(false, "failed to remove candidate {:?}", why);
+						}
+						None
+					}
+				})
+				.collect::<Vec<_>>()
+				.try_into()
+				.expect("filter_map operation can't result in a bounded vec larger than its original; qed")
+		}
+	}
+
+	/// Keep track of number of authored blocks per authority, uncles are counted as well since
+	/// they're a valid proof of being online.
+	impl<T: Config + pallet_authorship::Config>
+		pallet_authorship::EventHandler<T::AccountId, T::BlockNumber> for Pallet<T>
+	{
+		fn note_author(author: T::AccountId) {
+			let pot = Self::account_id();
+			// assumes an ED will be sent to pot.
+			let reward = T::Currency::free_balance(&pot)
+				.checked_sub(&T::Currency::minimum_balance())
+				.unwrap_or_else(Zero::zero)
+				.div(2u32.into());
+			// `reward` is half of pot account minus ED, this should never fail.
+			let _success = T::Currency::transfer(&pot, &author, reward, KeepAlive);
+			debug_assert!(_success.is_ok());
+			<LastAuthoredBlock<T>>::insert(author, frame_system::Pallet::<T>::block_number());
+
+			frame_system::Pallet::<T>::register_extra_weight_unchecked(
+				T::WeightInfo::note_author(),
+				DispatchClass::Mandatory,
+			);
+		}
+
+		fn note_uncle(_author: T::AccountId, _age: T::BlockNumber) {
+			//TODO can we ignore this?
+		}
+	}
+
+	/// Play the role of the session manager.
+	impl<T: Config> SessionManager<T::AccountId> for Pallet<T> {
+		fn new_session(index: SessionIndex) -> Option<Vec<T::AccountId>> {
+			log::info!(
+				"assembling new collators for new session {} at #{:?}",
+				index,
+				<frame_system::Pallet<T>>::block_number(),
+			);
+
+			let candidates = Self::candidates();
+			let candidates_len_before = candidates.len();
+			let active_candidates = Self::kick_stale_candidates(candidates);
+			let removed = candidates_len_before - active_candidates.len();
+			let result = Self::assemble_collators(active_candidates);
+
+			frame_system::Pallet::<T>::register_extra_weight_unchecked(
+				T::WeightInfo::new_session(candidates_len_before as u32, removed as u32),
+				DispatchClass::Mandatory,
+			);
+			Some(result)
+		}
+		fn start_session(_: SessionIndex) {
+			// we don't care.
+		}
+		fn end_session(_: SessionIndex) {
+			// we don't care.
+		}
+	}
+}

--- /dev/null
+++ b/pallets/collator-selection/src/mock.rs
@@ -0,0 +1,272 @@
+// Copyright 2019-2022 Unique Network (Gibraltar) Ltd.
+// This file is part of Unique Network.
+
+// Unique Network is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Unique Network is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Unique Network. If not, see <http://www.gnu.org/licenses/>.
+
+// Original license:
+// Copyright (C) 2021 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use super::*;
+use crate as collator_selection;
+use frame_support::{
+	ord_parameter_types, parameter_types,
+	traits::{FindAuthor, GenesisBuild, ValidatorRegistration},
+	PalletId,
+};
+use frame_system as system;
+use frame_system::EnsureSignedBy;
+use sp_core::H256;
+use sp_runtime::{
+	testing::{Header, UintAuthorityId},
+	traits::{BlakeTwo256, IdentityLookup, OpaqueKeys},
+	RuntimeAppPublic,
+};
+
+type UncheckedExtrinsic = frame_system::mocking::MockUncheckedExtrinsic<Test>;
+type Block = frame_system::mocking::MockBlock<Test>;
+
+// Configure a mock runtime to test the pallet.
+frame_support::construct_runtime!(
+	pub enum Test where
+		Block = Block,
+		NodeBlock = Block,
+		UncheckedExtrinsic = UncheckedExtrinsic,
+	{
+		System: frame_system::{Pallet, Call, Config, Storage, Event<T>},
+		Timestamp: pallet_timestamp::{Pallet, Call, Storage, Inherent},
+		Session: pallet_session::{Pallet, Call, Storage, Event, Config<T>},
+		Aura: pallet_aura::{Pallet, Storage, Config<T>},
+		Balances: pallet_balances::{Pallet, Call, Storage, Config<T>, Event<T>},
+		CollatorSelection: collator_selection::{Pallet, Call, Storage, Event<T>},
+		Authorship: pallet_authorship::{Pallet, Call, Storage, Inherent},
+	}
+);
+
+parameter_types! {
+	pub const BlockHashCount: u64 = 250;
+	pub const SS58Prefix: u8 = 42;
+}
+
+impl system::Config for Test {
+	type BaseCallFilter = frame_support::traits::Everything;
+	type BlockWeights = ();
+	type BlockLength = ();
+	type DbWeight = ();
+	type RuntimeOrigin = RuntimeOrigin;
+	type RuntimeCall = RuntimeCall;
+	type Index = u64;
+	type BlockNumber = u64;
+	type Hash = H256;
+	type Hashing = BlakeTwo256;
+	type AccountId = u64;
+	type Lookup = IdentityLookup<Self::AccountId>;
+	type Header = Header;
+	type RuntimeEvent = RuntimeEvent;
+	type BlockHashCount = BlockHashCount;
+	type Version = ();
+	type PalletInfo = PalletInfo;
+	type AccountData = pallet_balances::AccountData<u64>;
+	type OnNewAccount = ();
+	type OnKilledAccount = ();
+	type SystemWeightInfo = ();
+	type SS58Prefix = SS58Prefix;
+	type OnSetCode = ();
+	type MaxConsumers = frame_support::traits::ConstU32<16>;
+}
+
+parameter_types! {
+	pub const ExistentialDeposit: u64 = 5;
+	pub const MaxReserves: u32 = 50;
+}
+
+impl pallet_balances::Config for Test {
+	type Balance = u64;
+	type RuntimeEvent = RuntimeEvent;
+	type DustRemoval = ();
+	type ExistentialDeposit = ExistentialDeposit;
+	type AccountStore = System;
+	type WeightInfo = ();
+	type MaxLocks = ();
+	type MaxReserves = MaxReserves;
+	type ReserveIdentifier = [u8; 8];
+}
+
+pub struct Author4;
+impl FindAuthor<u64> for Author4 {
+	fn find_author<'a, I>(_digests: I) -> Option<u64>
+	where
+		I: 'a + IntoIterator<Item = (frame_support::ConsensusEngineId, &'a [u8])>,
+	{
+		Some(4)
+	}
+}
+
+impl pallet_authorship::Config for Test {
+	type FindAuthor = Author4;
+	type UncleGenerations = ();
+	type FilterUncle = ();
+	type EventHandler = CollatorSelection;
+}
+
+parameter_types! {
+	pub const MinimumPeriod: u64 = 1;
+}
+
+impl pallet_timestamp::Config for Test {
+	type Moment = u64;
+	type OnTimestampSet = Aura;
+	type MinimumPeriod = MinimumPeriod;
+	type WeightInfo = ();
+}
+
+impl pallet_aura::Config for Test {
+	type AuthorityId = sp_consensus_aura::sr25519::AuthorityId;
+	type MaxAuthorities = MaxAuthorities;
+	type DisabledValidators = ();
+}
+
+sp_runtime::impl_opaque_keys! {
+	pub struct MockSessionKeys {
+		// a key for aura authoring
+		pub aura: UintAuthorityId,
+	}
+}
+
+impl From<UintAuthorityId> for MockSessionKeys {
+	fn from(aura: sp_runtime::testing::UintAuthorityId) -> Self {
+		Self { aura }
+	}
+}
+
+parameter_types! {
+	pub static SessionHandlerCollators: Vec<u64> = Vec::new();
+	pub static SessionChangeBlock: u64 = 0;
+}
+
+pub struct TestSessionHandler;
+impl pallet_session::SessionHandler<u64> for TestSessionHandler {
+	const KEY_TYPE_IDS: &'static [sp_runtime::KeyTypeId] = &[UintAuthorityId::ID];
+	fn on_genesis_session<Ks: OpaqueKeys>(keys: &[(u64, Ks)]) {
+		SessionHandlerCollators::set(keys.into_iter().map(|(a, _)| *a).collect::<Vec<_>>())
+	}
+	fn on_new_session<Ks: OpaqueKeys>(_: bool, keys: &[(u64, Ks)], _: &[(u64, Ks)]) {
+		SessionChangeBlock::set(System::block_number());
+		dbg!(keys.len());
+		SessionHandlerCollators::set(keys.into_iter().map(|(a, _)| *a).collect::<Vec<_>>())
+	}
+	fn on_before_session_ending() {}
+	fn on_disabled(_: u32) {}
+}
+
+parameter_types! {
+	pub const Offset: u64 = 0;
+	pub const Period: u64 = 10;
+}
+
+impl pallet_session::Config for Test {
+	type RuntimeEvent = RuntimeEvent;
+	type ValidatorId = <Self as frame_system::Config>::AccountId;
+	// we don't have stash and controller, thus we don't need the convert as well.
+	type ValidatorIdOf = IdentityCollator;
+	type ShouldEndSession = pallet_session::PeriodicSessions<Period, Offset>;
+	type NextSessionRotation = pallet_session::PeriodicSessions<Period, Offset>;
+	type SessionManager = CollatorSelection;
+	type SessionHandler = TestSessionHandler;
+	type Keys = MockSessionKeys;
+	type WeightInfo = ();
+}
+
+ord_parameter_types! {
+	pub const RootAccount: u64 = 777;
+}
+
+parameter_types! {
+	pub const PotId: PalletId = PalletId(*b"PotStake");
+	pub const MaxCandidates: u32 = 20;
+	pub const MaxInvulnerables: u32 = 20;
+	pub const MinCandidates: u32 = 1;
+	pub const MaxAuthorities: u32 = 100_000;
+}
+
+pub struct IsRegistered;
+impl ValidatorRegistration<u64> for IsRegistered {
+	fn is_registered(id: &u64) -> bool {
+		if *id == 7u64 {
+			false
+		} else {
+			true
+		}
+	}
+}
+
+impl Config for Test {
+	type RuntimeEvent = RuntimeEvent;
+	type Currency = Balances;
+	type UpdateOrigin = EnsureSignedBy<RootAccount, u64>;
+	type PotId = PotId;
+	type MaxCandidates = MaxCandidates;
+	type MinCandidates = MinCandidates;
+	type MaxInvulnerables = MaxInvulnerables;
+	type KickThreshold = Period;
+	type ValidatorId = <Self as frame_system::Config>::AccountId;
+	type ValidatorIdOf = IdentityCollator;
+	type ValidatorRegistration = IsRegistered;
+	type WeightInfo = ();
+}
+
+pub fn new_test_ext() -> sp_io::TestExternalities {
+	sp_tracing::try_init_simple();
+	let mut t = frame_system::GenesisConfig::default().build_storage::<Test>().unwrap();
+	let invulnerables = vec![1, 2];
+
+	let balances = vec![(1, 100), (2, 100), (3, 100), (4, 100), (5, 100)];
+	let keys = balances
+		.iter()
+		.map(|&(i, _)| (i, i, MockSessionKeys { aura: UintAuthorityId(i) }))
+		.collect::<Vec<_>>();
+	let collator_selection = collator_selection::GenesisConfig::<Test> {
+		desired_candidates: 2,
+		candidacy_bond: 10,
+		invulnerables,
+	};
+	let session = pallet_session::GenesisConfig::<Test> { keys };
+	pallet_balances::GenesisConfig::<Test> { balances }
+		.assimilate_storage(&mut t)
+		.unwrap();
+	// collator selection must be initialized before session.
+	collator_selection.assimilate_storage(&mut t).unwrap();
+	session.assimilate_storage(&mut t).unwrap();
+
+	t.into()
+}
+
+pub fn initialize_to_block(n: u64) {
+	for i in System::block_number() + 1..=n {
+		System::set_block_number(i);
+		<AllPalletsWithSystem as frame_support::traits::OnInitialize<u64>>::on_initialize(i);
+	}
+}

--- /dev/null
+++ b/pallets/collator-selection/src/tests.rs
@@ -0,0 +1,414 @@
+// Copyright 2019-2022 Unique Network (Gibraltar) Ltd.
+// This file is part of Unique Network.
+
+// Unique Network is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Unique Network is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Unique Network. If not, see <http://www.gnu.org/licenses/>.
+
+// Original license:
+// Copyright (C) 2021 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate as collator_selection;
+use crate::{mock::*, CandidateInfo, Error};
+use frame_support::{
+	assert_noop, assert_ok,
+	traits::{Currency, GenesisBuild, OnInitialize},
+};
+use pallet_balances::Error as BalancesError;
+use sp_runtime::traits::BadOrigin;
+
+#[test]
+fn basic_setup_works() {
+	new_test_ext().execute_with(|| {
+		assert_eq!(CollatorSelection::desired_candidates(), 2);
+		assert_eq!(CollatorSelection::candidacy_bond(), 10);
+
+		assert!(CollatorSelection::candidates().is_empty());
+		assert_eq!(CollatorSelection::invulnerables(), vec![1, 2]);
+	});
+}
+
+#[test]
+fn it_should_set_invulnerables() {
+	new_test_ext().execute_with(|| {
+		let new_set = vec![1, 2, 3, 4];
+		assert_ok!(CollatorSelection::set_invulnerables(
+			RuntimeOrigin::signed(RootAccount::get()),
+			new_set.clone()
+		));
+		assert_eq!(CollatorSelection::invulnerables(), new_set);
+
+		// cannot set with non-root.
+		assert_noop!(
+			CollatorSelection::set_invulnerables(RuntimeOrigin::signed(1), new_set.clone()),
+			BadOrigin
+		);
+
+		// cannot set invulnerables without associated validator keys
+		let invulnerables = vec![7];
+		assert_noop!(
+			CollatorSelection::set_invulnerables(
+				RuntimeOrigin::signed(RootAccount::get()),
+				invulnerables.clone()
+			),
+			Error::<Test>::ValidatorNotRegistered
+		);
+	});
+}
+
+#[test]
+fn set_desired_candidates_works() {
+	new_test_ext().execute_with(|| {
+		// given
+		assert_eq!(CollatorSelection::desired_candidates(), 2);
+
+		// can set
+		assert_ok!(CollatorSelection::set_desired_candidates(
+			RuntimeOrigin::signed(RootAccount::get()),
+			7
+		));
+		assert_eq!(CollatorSelection::desired_candidates(), 7);
+
+		// rejects bad origin
+		assert_noop!(
+			CollatorSelection::set_desired_candidates(RuntimeOrigin::signed(1), 8),
+			BadOrigin
+		);
+	});
+}
+
+#[test]
+fn set_candidacy_bond() {
+	new_test_ext().execute_with(|| {
+		// given
+		assert_eq!(CollatorSelection::candidacy_bond(), 10);
+
+		// can set
+		assert_ok!(CollatorSelection::set_candidacy_bond(
+			RuntimeOrigin::signed(RootAccount::get()),
+			7
+		));
+		assert_eq!(CollatorSelection::candidacy_bond(), 7);
+
+		// rejects bad origin.
+		assert_noop!(CollatorSelection::set_candidacy_bond(RuntimeOrigin::signed(1), 8), BadOrigin);
+	});
+}
+
+#[test]
+fn cannot_register_candidate_if_too_many() {
+	new_test_ext().execute_with(|| {
+		// reset desired candidates:
+		<crate::DesiredCandidates<Test>>::put(0);
+
+		// can't accept anyone anymore.
+		assert_noop!(
+			CollatorSelection::register_as_candidate(RuntimeOrigin::signed(3)),
+			Error::<Test>::TooManyCandidates,
+		);
+
+		// reset desired candidates:
+		<crate::DesiredCandidates<Test>>::put(1);
+		assert_ok!(CollatorSelection::register_as_candidate(RuntimeOrigin::signed(4)));
+
+		// but no more
+		assert_noop!(
+			CollatorSelection::register_as_candidate(RuntimeOrigin::signed(5)),
+			Error::<Test>::TooManyCandidates,
+		);
+	})
+}
+
+#[test]
+fn cannot_unregister_candidate_if_too_few() {
+	new_test_ext().execute_with(|| {
+		// reset desired candidates:
+		<crate::DesiredCandidates<Test>>::put(1);
+		assert_ok!(CollatorSelection::register_as_candidate(RuntimeOrigin::signed(4)));
+
+		// can not remove too few
+		assert_noop!(
+			CollatorSelection::leave_intent(RuntimeOrigin::signed(4)),
+			Error::<Test>::TooFewCandidates,
+		);
+	})
+}
+
+#[test]
+fn cannot_register_as_candidate_if_invulnerable() {
+	new_test_ext().execute_with(|| {
+		assert_eq!(CollatorSelection::invulnerables(), vec![1, 2]);
+
+		// can't 1 because it is invulnerable.
+		assert_noop!(
+			CollatorSelection::register_as_candidate(RuntimeOrigin::signed(1)),
+			Error::<Test>::AlreadyInvulnerable,
+		);
+	})
+}
+
+#[test]
+fn cannot_register_as_candidate_if_keys_not_registered() {
+	new_test_ext().execute_with(|| {
+		// can't 7 because keys not registered.
+		assert_noop!(
+			CollatorSelection::register_as_candidate(RuntimeOrigin::signed(7)),
+			Error::<Test>::ValidatorNotRegistered
+		);
+	})
+}
+
+#[test]
+fn cannot_register_dupe_candidate() {
+	new_test_ext().execute_with(|| {
+		// can add 3 as candidate
+		assert_ok!(CollatorSelection::register_as_candidate(RuntimeOrigin::signed(3)));
+		let addition = CandidateInfo { who: 3, deposit: 10 };
+		assert_eq!(CollatorSelection::candidates(), vec![addition]);
+		assert_eq!(CollatorSelection::last_authored_block(3), 10);
+		assert_eq!(Balances::free_balance(3), 90);
+
+		// but no more
+		assert_noop!(
+			CollatorSelection::register_as_candidate(RuntimeOrigin::signed(3)),
+			Error::<Test>::AlreadyCandidate,
+		);
+	})
+}
+
+#[test]
+fn cannot_register_as_candidate_if_poor() {
+	new_test_ext().execute_with(|| {
+		assert_eq!(Balances::free_balance(&3), 100);
+		assert_eq!(Balances::free_balance(&33), 0);
+
+		// works
+		assert_ok!(CollatorSelection::register_as_candidate(RuntimeOrigin::signed(3)));
+
+		// poor
+		assert_noop!(
+			CollatorSelection::register_as_candidate(RuntimeOrigin::signed(33)),
+			BalancesError::<Test>::InsufficientBalance,
+		);
+	});
+}
+
+#[test]
+fn register_as_candidate_works() {
+	new_test_ext().execute_with(|| {
+		// given
+		assert_eq!(CollatorSelection::desired_candidates(), 2);
+		assert_eq!(CollatorSelection::candidacy_bond(), 10);
+		assert_eq!(CollatorSelection::candidates(), Vec::new());
+		assert_eq!(CollatorSelection::invulnerables(), vec![1, 2]);
+
+		// take two endowed, non-invulnerables accounts.
+		assert_eq!(Balances::free_balance(&3), 100);
+		assert_eq!(Balances::free_balance(&4), 100);
+
+		assert_ok!(CollatorSelection::register_as_candidate(RuntimeOrigin::signed(3)));
+		assert_ok!(CollatorSelection::register_as_candidate(RuntimeOrigin::signed(4)));
+
+		assert_eq!(Balances::free_balance(&3), 90);
+		assert_eq!(Balances::free_balance(&4), 90);
+
+		assert_eq!(CollatorSelection::candidates().len(), 2);
+	});
+}
+
+#[test]
+fn leave_intent() {
+	new_test_ext().execute_with(|| {
+		// register a candidate.
+		assert_ok!(CollatorSelection::register_as_candidate(RuntimeOrigin::signed(3)));
+		assert_eq!(Balances::free_balance(3), 90);
+
+		// register too so can leave above min candidates
+		assert_ok!(CollatorSelection::register_as_candidate(RuntimeOrigin::signed(5)));
+		assert_eq!(Balances::free_balance(5), 90);
+
+		// cannot leave if not candidate.
+		assert_noop!(
+			CollatorSelection::leave_intent(RuntimeOrigin::signed(4)),
+			Error::<Test>::NotCandidate
+		);
+
+		// bond is returned
+		assert_ok!(CollatorSelection::leave_intent(RuntimeOrigin::signed(3)));
+		assert_eq!(Balances::free_balance(3), 100);
+		assert_eq!(CollatorSelection::last_authored_block(3), 0);
+	});
+}
+
+#[test]
+fn authorship_event_handler() {
+	new_test_ext().execute_with(|| {
+		// put 100 in the pot + 5 for ED
+		Balances::make_free_balance_be(&CollatorSelection::account_id(), 105);
+
+		// 4 is the default author.
+		assert_eq!(Balances::free_balance(4), 100);
+		assert_ok!(CollatorSelection::register_as_candidate(RuntimeOrigin::signed(4)));
+		// triggers `note_author`
+		Authorship::on_initialize(1);
+
+		let collator = CandidateInfo { who: 4, deposit: 10 };
+
+		assert_eq!(CollatorSelection::candidates(), vec![collator]);
+		assert_eq!(CollatorSelection::last_authored_block(4), 0);
+
+		// half of the pot goes to the collator who's the author (4 in tests).
+		assert_eq!(Balances::free_balance(4), 140);
+		// half + ED stays.
+		assert_eq!(Balances::free_balance(CollatorSelection::account_id()), 55);
+	});
+}
+
+#[test]
+fn fees_edgecases() {
+	new_test_ext().execute_with(|| {
+		// Nothing panics, no reward when no ED in balance
+		Authorship::on_initialize(1);
+		// put some money into the pot at ED
+		Balances::make_free_balance_be(&CollatorSelection::account_id(), 5);
+		// 4 is the default author.
+		assert_eq!(Balances::free_balance(4), 100);
+		assert_ok!(CollatorSelection::register_as_candidate(RuntimeOrigin::signed(4)));
+		// triggers `note_author`
+		Authorship::on_initialize(1);
+
+		let collator = CandidateInfo { who: 4, deposit: 10 };
+
+		assert_eq!(CollatorSelection::candidates(), vec![collator]);
+		assert_eq!(CollatorSelection::last_authored_block(4), 0);
+		// Nothing received
+		assert_eq!(Balances::free_balance(4), 90);
+		// all fee stays
+		assert_eq!(Balances::free_balance(CollatorSelection::account_id()), 5);
+	});
+}
+
+#[test]
+fn session_management_works() {
+	new_test_ext().execute_with(|| {
+		initialize_to_block(1);
+
+		assert_eq!(SessionChangeBlock::get(), 0);
+		assert_eq!(SessionHandlerCollators::get(), vec![1, 2]);
+
+		initialize_to_block(4);
+
+		assert_eq!(SessionChangeBlock::get(), 0);
+		assert_eq!(SessionHandlerCollators::get(), vec![1, 2]);
+
+		// add a new collator
+		assert_ok!(CollatorSelection::register_as_candidate(RuntimeOrigin::signed(3)));
+
+		// session won't see this.
+		assert_eq!(SessionHandlerCollators::get(), vec![1, 2]);
+		// but we have a new candidate.
+		assert_eq!(CollatorSelection::candidates().len(), 1);
+
+		initialize_to_block(10);
+		assert_eq!(SessionChangeBlock::get(), 10);
+		// pallet-session has 1 session delay; current validators are the same.
+		assert_eq!(Session::validators(), vec![1, 2]);
+		// queued ones are changed, and now we have 3.
+		assert_eq!(Session::queued_keys().len(), 3);
+		// session handlers (aura, et. al.) cannot see this yet.
+		assert_eq!(SessionHandlerCollators::get(), vec![1, 2]);
+
+		initialize_to_block(20);
+		assert_eq!(SessionChangeBlock::get(), 20);
+		// changed are now reflected to session handlers.
+		assert_eq!(SessionHandlerCollators::get(), vec![1, 2, 3]);
+	});
+}
+
+#[test]
+fn kick_mechanism() {
+	new_test_ext().execute_with(|| {
+		// add a new collator
+		assert_ok!(CollatorSelection::register_as_candidate(RuntimeOrigin::signed(3)));
+		assert_ok!(CollatorSelection::register_as_candidate(RuntimeOrigin::signed(4)));
+		initialize_to_block(10);
+		assert_eq!(CollatorSelection::candidates().len(), 2);
+		initialize_to_block(20);
+		assert_eq!(SessionChangeBlock::get(), 20);
+		// 4 authored this block, gets to stay 3 was kicked
+		assert_eq!(CollatorSelection::candidates().len(), 1);
+		// 3 will be kicked after 1 session delay
+		assert_eq!(SessionHandlerCollators::get(), vec![1, 2, 3, 4]);
+		let collator = CandidateInfo { who: 4, deposit: 10 };
+		assert_eq!(CollatorSelection::candidates(), vec![collator]);
+		assert_eq!(CollatorSelection::last_authored_block(4), 20);
+		initialize_to_block(30);
+		// 3 gets kicked after 1 session delay
+		assert_eq!(SessionHandlerCollators::get(), vec![1, 2, 4]);
+		// kicked collator gets funds back
+		assert_eq!(Balances::free_balance(3), 100);
+	});
+}
+
+#[test]
+fn should_not_kick_mechanism_too_few() {
+	new_test_ext().execute_with(|| {
+		// add a new collator
+		assert_ok!(CollatorSelection::register_as_candidate(RuntimeOrigin::signed(3)));
+		assert_ok!(CollatorSelection::register_as_candidate(RuntimeOrigin::signed(5)));
+		initialize_to_block(10);
+		assert_eq!(CollatorSelection::candidates().len(), 2);
+		initialize_to_block(20);
+		assert_eq!(SessionChangeBlock::get(), 20);
+		// 4 authored this block, 5 gets to stay too few 3 was kicked
+		assert_eq!(CollatorSelection::candidates().len(), 1);
+		// 3 will be kicked after 1 session delay
+		assert_eq!(SessionHandlerCollators::get(), vec![1, 2, 3, 5]);
+		let collator = CandidateInfo { who: 5, deposit: 10 };
+		assert_eq!(CollatorSelection::candidates(), vec![collator]);
+		assert_eq!(CollatorSelection::last_authored_block(4), 20);
+		initialize_to_block(30);
+		// 3 gets kicked after 1 session delay
+		assert_eq!(SessionHandlerCollators::get(), vec![1, 2, 5]);
+		// kicked collator gets funds back
+		assert_eq!(Balances::free_balance(3), 100);
+	});
+}
+
+#[test]
+#[should_panic = "duplicate invulnerables in genesis."]
+fn cannot_set_genesis_value_twice() {
+	sp_tracing::try_init_simple();
+	let mut t = frame_system::GenesisConfig::default().build_storage::<Test>().unwrap();
+	let invulnerables = vec![1, 1];
+
+	let collator_selection = collator_selection::GenesisConfig::<Test> {
+		desired_candidates: 2,
+		candidacy_bond: 10,
+		invulnerables,
+	};
+	// collator selection must be initialized before session.
+	collator_selection.assimilate_storage(&mut t).unwrap();
+}

--- /dev/null
+++ b/pallets/collator-selection/src/weights.rs
@@ -0,0 +1,148 @@
+// Copyright 2019-2022 Unique Network (Gibraltar) Ltd.
+// This file is part of Unique Network.
+
+// Unique Network is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Unique Network is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Unique Network. If not, see <http://www.gnu.org/licenses/>.
+
+// Original license:
+// This file is part of Substrate.
+
+// Copyright (C) 2021 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#![allow(unused_parens)]
+#![allow(unused_imports)]
+
+use frame_support::{
+	traits::Get,
+	weights::{constants::RocksDbWeight, Weight},
+};
+use sp_std::marker::PhantomData;
+
+// The weight info trait for `pallet_collator_selection`.
+pub trait WeightInfo {
+	fn set_invulnerables(_b: u32) -> Weight;
+	fn set_desired_candidates() -> Weight;
+	fn set_candidacy_bond() -> Weight;
+	fn register_as_candidate(_c: u32) -> Weight;
+	fn leave_intent(_c: u32) -> Weight;
+	fn note_author() -> Weight;
+	fn new_session(_c: u32, _r: u32) -> Weight;
+}
+
+/// Weights for pallet_collator_selection using the Substrate node and recommended hardware.
+pub struct SubstrateWeight<T>(PhantomData<T>);
+impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {
+	fn set_invulnerables(b: u32) -> Weight {
+		Weight::from_ref_time(18_563_000 as u64)
+			// Standard Error: 0
+			.saturating_add(Weight::from_ref_time(68_000 as u64).saturating_mul(b as u64))
+			.saturating_add(T::DbWeight::get().writes(1 as u64))
+	}
+	fn set_desired_candidates() -> Weight {
+		Weight::from_ref_time(16_363_000 as u64).saturating_add(T::DbWeight::get().writes(1 as u64))
+	}
+	fn set_candidacy_bond() -> Weight {
+		Weight::from_ref_time(16_840_000 as u64).saturating_add(T::DbWeight::get().writes(1 as u64))
+	}
+	fn register_as_candidate(c: u32) -> Weight {
+		Weight::from_ref_time(71_196_000 as u64)
+			// Standard Error: 0
+			.saturating_add(Weight::from_ref_time(198_000 as u64).saturating_mul(c as u64))
+			.saturating_add(T::DbWeight::get().reads(4 as u64))
+			.saturating_add(T::DbWeight::get().writes(2 as u64))
+	}
+	fn leave_intent(c: u32) -> Weight {
+		Weight::from_ref_time(55_336_000 as u64)
+			// Standard Error: 0
+			.saturating_add(Weight::from_ref_time(151_000 as u64).saturating_mul(c as u64))
+			.saturating_add(T::DbWeight::get().reads(1 as u64))
+			.saturating_add(T::DbWeight::get().writes(2 as u64))
+	}
+	fn note_author() -> Weight {
+		Weight::from_ref_time(71_461_000 as u64)
+			.saturating_add(T::DbWeight::get().reads(3 as u64))
+			.saturating_add(T::DbWeight::get().writes(4 as u64))
+	}
+	fn new_session(r: u32, c: u32) -> Weight {
+		Weight::from_ref_time(0 as u64)
+			// Standard Error: 1_010_000
+			.saturating_add(Weight::from_ref_time(109_961_000 as u64).saturating_mul(r as u64))
+			// Standard Error: 1_010_000
+			.saturating_add(Weight::from_ref_time(151_952_000 as u64).saturating_mul(c as u64))
+			.saturating_add(T::DbWeight::get().reads((1 as u64).saturating_mul(r as u64)))
+			.saturating_add(T::DbWeight::get().reads((2 as u64).saturating_mul(c as u64)))
+			.saturating_add(T::DbWeight::get().writes((2 as u64).saturating_mul(r as u64)))
+			.saturating_add(T::DbWeight::get().writes((2 as u64).saturating_mul(c as u64)))
+	}
+}
+
+// For backwards compatibility and tests
+impl WeightInfo for () {
+	fn set_invulnerables(b: u32) -> Weight {
+		Weight::from_ref_time(18_563_000 as u64)
+			// Standard Error: 0
+			.saturating_add(Weight::from_ref_time(68_000 as u64).saturating_mul(b as u64))
+			.saturating_add(RocksDbWeight::get().writes(1 as u64))
+	}
+	fn set_desired_candidates() -> Weight {
+		Weight::from_ref_time(16_363_000 as u64)
+			.saturating_add(RocksDbWeight::get().writes(1 as u64))
+	}
+	fn set_candidacy_bond() -> Weight {
+		Weight::from_ref_time(16_840_000 as u64)
+			.saturating_add(RocksDbWeight::get().writes(1 as u64))
+	}
+	fn register_as_candidate(c: u32) -> Weight {
+		Weight::from_ref_time(71_196_000 as u64)
+			// Standard Error: 0
+			.saturating_add(Weight::from_ref_time(198_000 as u64).saturating_mul(c as u64))
+			.saturating_add(RocksDbWeight::get().reads(4 as u64))
+			.saturating_add(RocksDbWeight::get().writes(2 as u64))
+	}
+	fn leave_intent(c: u32) -> Weight {
+		Weight::from_ref_time(55_336_000 as u64)
+			// Standard Error: 0
+			.saturating_add(Weight::from_ref_time(151_000 as u64).saturating_mul(c as u64))
+			.saturating_add(RocksDbWeight::get().reads(1 as u64))
+			.saturating_add(RocksDbWeight::get().writes(2 as u64))
+	}
+	fn note_author() -> Weight {
+		Weight::from_ref_time(71_461_000 as u64)
+			.saturating_add(RocksDbWeight::get().reads(3 as u64))
+			.saturating_add(RocksDbWeight::get().writes(4 as u64))
+	}
+	fn new_session(r: u32, c: u32) -> Weight {
+		Weight::from_ref_time(0 as u64)
+			// Standard Error: 1_010_000
+			.saturating_add(Weight::from_ref_time(109_961_000 as u64).saturating_mul(r as u64))
+			// Standard Error: 1_010_000
+			.saturating_add(Weight::from_ref_time(151_952_000 as u64).saturating_mul(c as u64))
+			.saturating_add(RocksDbWeight::get().reads((1 as u64).saturating_mul(r as u64)))
+			.saturating_add(RocksDbWeight::get().reads((2 as u64).saturating_mul(c as u64)))
+			.saturating_add(RocksDbWeight::get().writes((2 as u64).saturating_mul(r as u64)))
+			.saturating_add(RocksDbWeight::get().writes((2 as u64).saturating_mul(c as u64)))
+	}
+}