clear votes with tests

ImbueNetwork · Dec 8, 2023 · 8cfee80 · 8cfee80
1 parent 86225b7
commit 8cfee80
Show file tree

Hide file tree

Showing 2 changed files with 70 additions and 1 deletion.
diff --git a/pallets/proposals/src/impls/pallet_impls.rs b/pallets/proposals/src/impls/pallet_impls.rs
@@ -316,7 +316,13 @@ impl<T: Config> Pallet<T> {
         let exp_block = Rounds::<T>::take((project_key, milestone_key), RoundType::VotingRound)
             .ok_or(Error::<T>::VotingRoundNotStarted)?;
         // Prevent hook from calling.
+        // TODO: only remove project key????
         RoundsExpiring::<T>::remove(exp_block);
+
+        MilestoneVotes::<T>::mutate(project_key, |btree_votes|{
+          let _val = btree_votes.remove(&milestone_key);
+        });
+
         // Allow future votes to occur on this milestone
         IndividualVoteStore::<T>::try_mutate(project_key, |maybe_individual_votes| {
             if let Some(individual_votes) = maybe_individual_votes {

diff --git a/pallets/proposals/src/tests/pallet.rs b/pallets/proposals/src/tests/pallet.rs
@@ -541,7 +541,7 @@ fn vote_on_milestone_not_contributor() {
 #[test]
 fn vote_on_milestone_actually_adds_to_vote() {
     build_test_externality().execute_with(|| {
-        let cont = get_contributions::<Test>(vec![BOB, CHARLIE], 100_000);
+        let cont = get_contributions::<Test>(vec![BOB, CHARLIE, DAVE], 100_000);
         let prop_milestones = get_milestones(10);
         let project_key =
             create_and_fund_project::<Test>(ALICE, cont, prop_milestones, CurrencyId::Native)
@@ -615,6 +615,69 @@ fn vote_on_milestone_autofinalises_on_all_voted_and_fail() {
     });
 }
 
+fn vote_struct_removed_on_autofinalisation_success() {
+    build_test_externality().execute_with(|| {
+        let cont = get_contributions::<Test>(vec![BOB, CHARLIE], 100_000);
+        let prop_milestones = get_milestones(10);
+        let project_key =
+            create_and_fund_project::<Test>(ALICE, cont, prop_milestones, CurrencyId::Native)
+                .unwrap();
+        let milestone_key = 5;
+        assert_ok!(Proposals::submit_milestone(
+            RuntimeOrigin::signed(ALICE),
+            project_key,
+            milestone_key
+        ));
+        assert_ok!(Proposals::vote_on_milestone(
+            RuntimeOrigin::signed(BOB),
+            project_key,
+            milestone_key,
+            true
+        ));
+        assert!(MilestoneVotes::<Test>::get(project_key).get(&milestone_key).is_some());
+
+        assert_ok!(Proposals::vote_on_milestone(
+            RuntimeOrigin::signed(CHARLIE),
+            project_key,
+            milestone_key,
+            true
+        ));
+        assert!(MilestoneVotes::<Test>::get(project_key).get(&milestone_key).is_none());
+    });
+}
+
+fn vote_struct_removed_on_autofinalisation_failure() {
+    build_test_externality().execute_with(|| {
+        let cont = get_contributions::<Test>(vec![BOB, CHARLIE], 100_000);
+        let prop_milestones = get_milestones(10);
+        let project_key =
+            create_and_fund_project::<Test>(ALICE, cont, prop_milestones, CurrencyId::Native)
+                .unwrap();
+        let milestone_key = 5;
+        assert_ok!(Proposals::submit_milestone(
+            RuntimeOrigin::signed(ALICE),
+            project_key,
+            milestone_key
+        ));
+        assert_ok!(Proposals::vote_on_milestone(
+            RuntimeOrigin::signed(BOB),
+            project_key,
+            milestone_key,
+            false
+        ));
+        assert!(MilestoneVotes::<Test>::get(project_key).get(&milestone_key).is_some());
+        assert_ok!(Proposals::vote_on_milestone(
+            RuntimeOrigin::signed(CHARLIE),
+            project_key,
+            milestone_key,
+            false
+        ));
+
+        assert!(MilestoneVotes::<Test>::get(project_key).get(&milestone_key).is_none());
+    });
+}
+
+
 #[test]
 fn withdraw_not_initiator() {
     build_test_externality().execute_with(|| {