Society: Ensure all votes are removed after tally (#4666)

* Ensure all votes are removed after tally * Fix comment
paritytech · Jan 17, 2020 · 8de5340 · 8de5340
1 parent f0c1852
commit 8de5340
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 3 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/bin/node/cli/Cargo.toml b/bin/node/cli/Cargo.toml
@@ -72,7 +72,6 @@ pallet-transaction-payment = { version = "2.0.0", path = "../../../frame/transac
 frame-support = { version = "2.0.0", default-features = false, path = "../../../frame/support" }
 pallet-im-online = { version = "2.0.0", default-features = false, path = "../../../frame/im-online" }
 pallet-authority-discovery = { version = "2.0.0",  path = "../../../frame/authority-discovery" }
-pallet-society = { version = "2.0.0",  path = "../../../frame/society" }
 
 # node-specific dependencies
 node-runtime = { version = "2.0.0", path = "../runtime" }

diff --git a/frame/society/src/lib.rs b/frame/society/src/lib.rs
@@ -1324,6 +1324,9 @@ impl<T: Trait<I>, I: Instance> Module<T, I> {
 				}
 			}).collect::<Vec<_>>();
 
+			// Clean up all votes.
+			<Votes<T, I>>::remove_all();
+
 			// Reward one of the voters who voted the right way.
 			if !total_slash.is_zero() {
 				if let Some(winner) = pick_item(&mut rng, &rewardees) {
@@ -1472,7 +1475,7 @@ impl<T: Trait<I>, I: Instance> Module<T, I> {
 				let mut rejection_count = 0;
 				// Tallies total number of approve and reject votes for the defender.
 				members.iter()
-					.filter_map(|m| <DefenderVotes<T, I>>::get(m))
+					.filter_map(|m| <DefenderVotes<T, I>>::take(m))
 					.for_each(|v| {
 						match v {
 							Vote::Approve => approval_count += 1,
@@ -1485,6 +1488,9 @@ impl<T: Trait<I>, I: Instance> Module<T, I> {
 					Self::suspend_member(&defender);
 					*members = Self::members();
 				}
+
+				// Clean up all votes.
+				<DefenderVotes<T, I>>::remove_all();
 			}
 
 			// Start a new defender rotation

diff --git a/frame/society/src/tests.rs b/frame/society/src/tests.rs
@@ -538,6 +538,11 @@ fn challenges_work() {
 		assert_ok!(Society::add_member(&20));
 		assert_ok!(Society::add_member(&30));
 		assert_ok!(Society::add_member(&40));
+		// Votes are empty
+		assert_eq!(<DefenderVotes<Test>>::get(10), None);
+		assert_eq!(<DefenderVotes<Test>>::get(20), None);
+		assert_eq!(<DefenderVotes<Test>>::get(30), None);
+		assert_eq!(<DefenderVotes<Test>>::get(40), None);
 		// Check starting point
 		assert_eq!(Society::members(), vec![10, 20, 30, 40]);
 		assert_eq!(Society::defender(), None);
@@ -561,6 +566,11 @@ fn challenges_work() {
 		run_to_block(24);
 		// 20 survives
 		assert_eq!(Society::members(), vec![10, 20, 30, 40]);
+		// Votes are reset
+		assert_eq!(<DefenderVotes<Test>>::get(10), None);
+		assert_eq!(<DefenderVotes<Test>>::get(20), None);
+		assert_eq!(<DefenderVotes<Test>>::get(30), None);
+		assert_eq!(<DefenderVotes<Test>>::get(40), None);
 		// One more time
 		assert_eq!(Society::defender(), Some(20));
 		// 2 people say accept, 2 reject
@@ -574,6 +584,11 @@ fn challenges_work() {
 		assert_eq!(Society::suspended_member(20), true);
 		// New defender is chosen
 		assert_eq!(Society::defender(), Some(40));
+		// Votes are reset
+		assert_eq!(<DefenderVotes<Test>>::get(10), None);
+		assert_eq!(<DefenderVotes<Test>>::get(20), None);
+		assert_eq!(<DefenderVotes<Test>>::get(30), None);
+		assert_eq!(<DefenderVotes<Test>>::get(40), None);
 	});
 }