Fix vtoken voting (#1540)

* Separate the logic for retrieving the current block for different tokens. * change ReferendumTimeout max number to 100 * change RelaychainBlockNumber to agent block number --------- Co-authored-by: Edwin <[email protected]>
bifrost-io · Dec 4, 2024 · 26119e5 · 26119e5
1 parent 6d2ab67
commit 26119e5
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 16 deletions.
diff --git a/pallets/vtoken-voting/src/lib.rs b/pallets/vtoken-voting/src/lib.rs
@@ -1150,7 +1150,8 @@ pub mod pallet {
 										.ok_or(Error::<T>::NoData)?
 										.saturating_mul(lock_periods.into()),
 								);
-								let now = T::RelaychainBlockNumberProvider::current_block_number();
+
+								let now = Self::get_agent_block_number(&vtoken)?;
 								if now < unlock_at {
 									ensure!(
 										matches!(scope, UnvoteScope::Any),
@@ -1174,8 +1175,9 @@ pub mod pallet {
 		/// Rejig the lock on an account. It will never get more stringent (since that would
 		/// indicate a security hole) but may be reduced from what they are currently.
 		pub(crate) fn update_lock(who: &AccountIdOf<T>, vtoken: CurrencyIdOf<T>) -> DispatchResult {
+			let current_block = Self::get_agent_block_number(&vtoken)?;
 			let lock_needed = VotingFor::<T>::mutate(who, |voting| {
-				voting.rejig(T::RelaychainBlockNumberProvider::current_block_number());
+				voting.rejig(current_block);
 				voting.locked_balance()
 			});
 
@@ -1278,8 +1280,10 @@ pub mod pallet {
 				(Some(ReferendumInfo::Completed(moment)), Some((lock_periods, _balance))) => {
 					let locking_period =
 						VoteLockingPeriod::<T>::get(vtoken).ok_or(Error::<T>::NoData)?;
+
+					let current_block = Self::get_agent_block_number(&vtoken)?;
 					ensure!(
-						T::RelaychainBlockNumberProvider::current_block_number() >=
+						current_block >=
 							moment.saturating_add(
 								locking_period.saturating_mul(lock_periods.into())
 							),
@@ -1288,10 +1292,8 @@ pub mod pallet {
 					Ok(())
 				},
 				(Some(ReferendumInfo::Completed(moment)), None) => {
-					ensure!(
-						T::RelaychainBlockNumberProvider::current_block_number() >= moment,
-						Error::<T>::NotExpired
-					);
+					let current_block = Self::get_agent_block_number(&vtoken)?;
+					ensure!(current_block >= moment, Error::<T>::NotExpired);
 					Ok(())
 				},
 				_ => Err(Error::<T>::NotExpired.into()),

diff --git a/pallets/vtoken-voting/src/tests/vbnc_test.rs b/pallets/vtoken-voting/src/tests/vbnc_test.rs
@@ -143,7 +143,7 @@ fn successful_but_zero_conviction_vote_balance_can_be_unlocked() {
 			assert_ok!(VtokenVoting::unlock(RuntimeOrigin::signed(BOB), vtoken, poll_index));
 			assert_eq!(usable_balance(vtoken, &BOB), 20);
 
-			RelaychainDataProvider::set_block_number(13);
+			System::set_block_number(13);
 			assert_ok!(VtokenVoting::unlock(RuntimeOrigin::signed(ALICE), vtoken, poll_index));
 			assert_eq!(usable_balance(vtoken, &ALICE), 10);
 		});
@@ -181,7 +181,7 @@ fn unsuccessful_conviction_vote_balance_can_be_unlocked() {
 				poll_index,
 				ReferendumInfoOf::<Runtime>::Completed(3),
 			));
-			RelaychainDataProvider::set_block_number(13);
+			System::set_block_number(13);
 			assert_ok!(VtokenVoting::try_remove_vote(&ALICE, vtoken, poll_index, UnvoteScope::Any));
 			assert_ok!(VtokenVoting::update_lock(&ALICE, vtoken));
 			assert_eq!(usable_balance(vtoken, &ALICE), 10);
@@ -221,7 +221,7 @@ fn ensure_balance_after_unlock() {
 				poll_index,
 				ReferendumInfoOf::<Runtime>::Completed(3),
 			));
-			RelaychainDataProvider::set_block_number(13);
+			System::set_block_number(13);
 			assert_ok!(VtokenVoting::unlock(RuntimeOrigin::signed(ALICE), vtoken, poll_index));
 			assert_eq!(usable_balance(vtoken, &ALICE), 0);
 			assert_eq!(Tokens::accounts(&ALICE, vtoken).frozen, 10);
@@ -269,7 +269,7 @@ fn ensure_comprehensive_balance_after_unlock() {
 				poll_index,
 				ReferendumInfoOf::<Runtime>::Completed(3),
 			));
-			RelaychainDataProvider::set_block_number(13);
+			System::set_block_number(13);
 			assert_ok!(VtokenVoting::unlock(RuntimeOrigin::signed(ALICE), vtoken, poll_index));
 			assert_eq!(usable_balance(vtoken, &ALICE), 8);
 			assert_eq!(Tokens::accounts(&ALICE, vtoken).frozen, 2);
@@ -314,7 +314,7 @@ fn successful_conviction_vote_balance_stays_locked_for_correct_time() {
 				poll_index,
 				ReferendumInfoOf::<Runtime>::Completed(3),
 			));
-			RelaychainDataProvider::set_block_number(163);
+			System::set_block_number(163);
 			for i in 1..=5 {
 				assert_ok!(VtokenVoting::try_remove_vote(&i, vtoken, poll_index, UnvoteScope::Any));
 			}
@@ -392,15 +392,15 @@ fn lock_amalgamation_valid_with_multiple_removed_votes() {
 					.unwrap()
 			);
 
-			RelaychainDataProvider::set_block_number(10);
+			System::set_block_number(10);
 			assert_noop!(
 				VtokenVoting::unlock(RuntimeOrigin::signed(ALICE), vtoken, 0),
 				Error::<Runtime>::NoPermissionYet
 			);
 			assert_eq!(VotingFor::<Runtime>::get(&ALICE).locked_balance(), 10);
 			assert_eq!(usable_balance(vtoken, &ALICE), 0);
 
-			RelaychainDataProvider::set_block_number(11);
+			System::set_block_number(11);
 			assert_ok!(VtokenVoting::unlock(RuntimeOrigin::signed(ALICE), vtoken, 0));
 			assert_eq!(VotingFor::<Runtime>::get(&ALICE).locked_balance(), 10);
 			assert_eq!(usable_balance(vtoken, &ALICE), 0);
@@ -410,7 +410,7 @@ fn lock_amalgamation_valid_with_multiple_removed_votes() {
 					.unwrap()
 			);
 
-			RelaychainDataProvider::set_block_number(11);
+			System::set_block_number(11);
 			assert_ok!(VtokenVoting::unlock(RuntimeOrigin::signed(ALICE), vtoken, 1));
 			assert_eq!(usable_balance(vtoken, &ALICE), 0);
 			assert_eq!(
@@ -419,7 +419,7 @@ fn lock_amalgamation_valid_with_multiple_removed_votes() {
 					.unwrap()
 			);
 
-			RelaychainDataProvider::set_block_number(21);
+			System::set_block_number(21);
 			assert_ok!(VtokenVoting::unlock(RuntimeOrigin::signed(ALICE), vtoken, 2));
 			assert_eq!(usable_balance(vtoken, &ALICE), 10);
 			assert_eq!(