From 216d712b42fda273fa3c436d1e5c4766329630a7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gon=C3=A7alo=20Pestana?= <g6pestana@gmail.com>
Date: Thu, 14 Nov 2024 16:27:32 +0100
Subject: [PATCH] adds epm tests for max backers per winner

---
 .../election-provider-multi-phase/src/mock.rs |  4 +
 .../src/unsigned.rs                           | 92 +++++++++++++++++++
 2 files changed, 96 insertions(+)

diff --git a/substrate/frame/election-provider-multi-phase/src/mock.rs b/substrate/frame/election-provider-multi-phase/src/mock.rs
index 89c99c8f94a4..100abc01c2d9 100644
--- a/substrate/frame/election-provider-multi-phase/src/mock.rs
+++ b/substrate/frame/election-provider-multi-phase/src/mock.rs
@@ -611,6 +611,10 @@ impl ExtBuilder {
 		<SignedMaxWeight>::set(weight);
 		self
 	}
+	pub fn max_backers_per_winner(self, max: u32) -> Self {
+		<MaxBackersPerWinner>::set(max);
+		self
+	}
 	pub fn build(self) -> sp_io::TestExternalities {
 		sp_tracing::try_init_simple();
 		let mut storage =
diff --git a/substrate/frame/election-provider-multi-phase/src/unsigned.rs b/substrate/frame/election-provider-multi-phase/src/unsigned.rs
index 7ce8a4a1d6af..82401967b8af 100644
--- a/substrate/frame/election-provider-multi-phase/src/unsigned.rs
+++ b/substrate/frame/election-provider-multi-phase/src/unsigned.rs
@@ -1865,6 +1865,98 @@ mod tests {
 		})
 	}
 
+	#[test]
+	fn mine_solution_always_respects_max_backers_per_winner() {
+		use crate::mock::MaxBackersPerWinner;
+		use frame_election_provider_support::BoundedSupport;
+
+		let targets = vec![10, 20, 30, 40];
+		let voters = vec![
+			(1, 10, bounded_vec![10, 20, 30]),
+			(2, 10, bounded_vec![10, 20, 30]),
+			(3, 10, bounded_vec![10, 20, 30]),
+			(4, 10, bounded_vec![10, 20, 30]),
+			(5, 10, bounded_vec![10, 20, 40]),
+		];
+		let snapshot = RoundSnapshot { voters: voters.clone(), targets: targets.clone() };
+		let (round, desired_targets) = (1, 3);
+
+		let score_unbounded =
+			ElectionScore { minimal_stake: 12, sum_stake: 50, sum_stake_squared: 874 };
+		let score_bounded =
+			ElectionScore { minimal_stake: 2, sum_stake: 10, sum_stake_squared: 44 };
+
+		// solution with max backers per winner will be lower.
+		assert!(score_unbounded > score_bounded);
+
+		// election with unbounded max backers per winnner.
+		ExtBuilder::default().max_backers_per_winner(u32::MAX).build_and_execute(|| {
+			assert_eq!(MaxBackersPerWinner::get(), u32::MAX);
+
+			let solution = Miner::<Runtime>::mine_solution_with_snapshot::<
+				<Runtime as Config>::Solver,
+			>(voters.clone(), targets.clone(), desired_targets)
+			.unwrap()
+			.0;
+
+			let ready_solution = Miner::<Runtime>::feasibility_check(
+				RawSolution { solution, score: score_unbounded, round },
+				Default::default(),
+				desired_targets,
+				snapshot.clone(),
+				round,
+				Default::default(),
+			)
+			.unwrap();
+
+			assert_eq!(
+				ready_solution.supports.into_iter().collect::<Vec<_>>(),
+				vec![
+					(
+						10,
+						BoundedSupport { total: 21, voters: bounded_vec![(1, 10), (4, 8), (5, 3)] }
+					),
+					(20, BoundedSupport { total: 17, voters: bounded_vec![(2, 10), (5, 7)] }),
+					(30, BoundedSupport { total: 12, voters: bounded_vec![(3, 10), (4, 2)] }),
+				]
+			);
+		});
+
+		// election with max 1 backer per winnner.
+		ExtBuilder::default().max_backers_per_winner(1).build_and_execute(|| {
+			assert_eq!(MaxBackersPerWinner::get(), 1);
+
+			let solution = Miner::<Runtime>::mine_solution_with_snapshot::<
+				<Runtime as Config>::Solver,
+			>(voters, targets, desired_targets)
+			.unwrap()
+			.0;
+
+			let ready_solution = Miner::<Runtime>::feasibility_check(
+				RawSolution { solution, score: score_bounded, round },
+				Default::default(),
+				desired_targets,
+				snapshot,
+				round,
+				Default::default(),
+			)
+			.unwrap();
+
+			for (_, supports) in ready_solution.supports.iter() {
+				assert!((supports.voters.len() as u32) <= MaxBackersPerWinner::get());
+			}
+
+			assert_eq!(
+				ready_solution.supports.into_iter().collect::<Vec<_>>(),
+				vec![
+					(10, BoundedSupport { total: 6, voters: bounded_vec![(1, 6)] }),
+					(20, BoundedSupport { total: 2, voters: bounded_vec![(1, 2)] }),
+					(30, BoundedSupport { total: 2, voters: bounded_vec![(1, 2)] }),
+				]
+			);
+		});
+	}
+
 	#[test]
 	fn trim_assignments_length_does_not_modify_when_short_enough() {
 		ExtBuilder::default().build_and_execute(|| {