More lenient mechanism for identifying stash accounts in purge_keys (#10004)

* More lenient StashOf finder

* Slightly safer version of the change

* Slightly safer version of the change

* Tests

* Formatting

substrate/frame/session/src/mock.rs

@@ -22,13 +22,13 @@ use crate as pallet_session;
 #[cfg(feature = "historical")]
 use crate::historical as pallet_session_historical;
 
-use std::cell::RefCell;
+use std::{cell::RefCell, collections::BTreeMap};
 
 use sp_core::{crypto::key_types::DUMMY, H256};
 use sp_runtime::{
 	impl_opaque_keys,
 	testing::{Header, UintAuthorityId},
-	traits::{BlakeTwo256, ConvertInto, IdentityLookup},
+	traits::{BlakeTwo256, IdentityLookup},
 };
 use sp_staking::SessionIndex;
 
@@ -111,6 +111,7 @@ thread_local! {
 	pub static DISABLED: RefCell<bool> = RefCell::new(false);
 	// Stores if `on_before_session_end` was called
 	pub static BEFORE_SESSION_END_CALLED: RefCell<bool> = RefCell::new(false);
+	pub static VALIDATOR_ACCOUNTS: RefCell<BTreeMap<u64, u64>> = RefCell::new(BTreeMap::new());
 }
 
 pub struct TestShouldEndSession;
@@ -225,6 +226,10 @@ pub fn new_test_ext() -> sp_io::TestExternalities {
 	pallet_session::GenesisConfig::<Test> { keys }
 		.assimilate_storage(&mut t)
 		.unwrap();
+	NEXT_VALIDATORS.with(|l| {
+		let v = l.borrow().iter().map(|&i| (i, i)).collect();
+		VALIDATOR_ACCOUNTS.with(|m| *m.borrow_mut() = v);
+	});
 	sp_io::TestExternalities::new(t)
 }
 
@@ -268,6 +273,18 @@ impl pallet_timestamp::Config for Test {
 	type WeightInfo = ();
 }
 
+pub struct TestValidatorIdOf;
+impl TestValidatorIdOf {
+	pub fn set(v: BTreeMap<u64, u64>) {
+		VALIDATOR_ACCOUNTS.with(|m| *m.borrow_mut() = v);
+	}
+}
+impl Convert<u64, Option<u64>> for TestValidatorIdOf {
+	fn convert(x: u64) -> Option<u64> {
+		VALIDATOR_ACCOUNTS.with(|m| m.borrow().get(&x).cloned())
+	}
+}
+
 impl Config for Test {
 	type ShouldEndSession = TestShouldEndSession;
 	#[cfg(feature = "historical")]
@@ -276,7 +293,7 @@ impl Config for Test {
 	type SessionManager = TestSessionManager;
 	type SessionHandler = TestSessionHandler;
 	type ValidatorId = u64;
-	type ValidatorIdOf = ConvertInto;
+	type ValidatorIdOf = TestValidatorIdOf;
 	type Keys = MockSessionKeys;
 	type Event = Event;
 	type NextSessionRotation = ();