use Option in more places, use TimestampProvider in PSP22Votes

Brushfam · Aug 16, 2023 · a34a8bf · a34a8bf
1 parent 1035a1c
commit a34a8bf
Show file tree

Hide file tree

Showing 10 changed files with 64 additions and 42 deletions.
diff --git a/contracts/src/governance/extensions/governor_quorum/impls.rs b/contracts/src/governance/extensions/governor_quorum/impls.rs
@@ -9,7 +9,7 @@ use crate::{
     governor::TimestampProvider,
     traits::{
         errors::GovernanceError,
-        governance::utils::VotesWrapper,
+        governance::utils::VotesRef,
     },
 };
 use openbrush::traits::{
@@ -63,7 +63,7 @@ pub trait QuorumImpl: Storage<Data> + Storage<governor_votes::Data> + QuorumEven
             .get()
             .ok_or(GovernanceError::TokenNotSet)?;
 
-        let past_total_supply = VotesWrapper::get_past_total_supply(&mut token, time_point)?;
+        let past_total_supply = VotesRef::get_past_total_supply(&mut token, time_point)?;
 
         past_total_supply
             .checked_mul(self.quorum_numerator_at(time_point))

diff --git a/contracts/src/governance/extensions/governor_votes/internal.rs b/contracts/src/governance/extensions/governor_votes/internal.rs
@@ -2,7 +2,7 @@ use crate::{
     extensions::governor_votes::Data,
     traits::{
         errors::GovernanceError,
-        governance::utils::VotesWrapper,
+        governance::utils::*,
     },
 };
 use ink::{
@@ -27,8 +27,8 @@ pub trait GovernorVotesInternal: Storage<Data> {
         timepoint: Timestamp,
         _params: Vec<u8>,
     ) -> Result<Balance, GovernanceError> {
-        let mut token = self.data().token.get().ok_or(GovernanceError::TokenNotSet)?;
+        let token = self.data().token.get().ok_or(GovernanceError::TokenNotSet)?;
 
-        VotesWrapper::get_past_votes(&mut token, account, timepoint)
+        VotesRef::get_past_votes(&token, account, timepoint)
     }
 }
diff --git a/contracts/src/governance/utils/votes/data.rs b/contracts/src/governance/utils/votes/data.rs
@@ -7,7 +7,7 @@ use openbrush::{
 #[derive(Default, Debug)]
 #[openbrush::storage_item]
 pub struct Data {
-    pub delegation: Mapping<AccountId, AccountId>,
+    pub delegation: Mapping<Option<AccountId>, AccountId>,
     pub delegate_checkpoints: Mapping<AccountId, Checkpoints>,
     #[lazy]
     pub total_checkpoints: Checkpoints,

diff --git a/contracts/src/governance/utils/votes/events.rs b/contracts/src/governance/utils/votes/events.rs
@@ -6,7 +6,7 @@ use openbrush::traits::{
 pub trait VotesEvents {
     fn emit_delegate_changed_event(
         &self,
-        _delegator: &AccountId,
+        _delegator: &Option<AccountId>,
         _from_delegate: &Option<AccountId>,
         _to_delegate: &Option<AccountId>,
     ) {

diff --git a/contracts/src/governance/utils/votes/impls.rs b/contracts/src/governance/utils/votes/impls.rs
@@ -63,12 +63,12 @@ pub trait VotesImpl: Storage<Data> + VotesInternal + NoncesImpl + VotesEvents +
     }
 
     fn delegates(&mut self, delegator: AccountId) -> Option<AccountId> {
-        self._delegates(&delegator)
+        self._delegates(&Some(delegator))
     }
 
     fn delegate(&mut self, delegatee: AccountId) -> Result<(), GovernanceError> {
         let account = Self::env().caller();
-        self._delegate(&account, &delegatee)
+        self._delegate(&Some(account), &Some(delegatee))
     }
 
     fn delegate_by_signature(
@@ -88,7 +88,7 @@ pub trait VotesImpl: Storage<Data> + VotesInternal + NoncesImpl + VotesEvents +
             return Err(GovernanceError::InvalidSignature(signer))
         } else {
             self._use_checked_nonce(&signer, nonce)?;
-            self._delegate(&signer, &delegatee)
+            self._delegate(&Some(signer), &Some(delegatee))
         }
     }
 }
diff --git a/contracts/src/governance/utils/votes/internal.rs b/contracts/src/governance/utils/votes/internal.rs
@@ -23,35 +23,41 @@ pub trait VotesInternal: Storage<Data> + VotesEvents {
         self.data::<Data>().total_checkpoints.get_or_default().latest()
     }
 
-    fn _delegates(&self, delegator: &AccountId) -> Option<AccountId> {
+    fn _delegates(&self, delegator: &Option<AccountId>) -> Option<AccountId> {
         self.data::<Data>().delegation.get(&delegator)
     }
 
-    fn _delegate(&mut self, delegator: &AccountId, delegatee: &AccountId) -> Result<(), GovernanceError> {
+    fn _delegate(
+        &mut self,
+        delegator: &Option<AccountId>,
+        delegatee: &Option<AccountId>,
+    ) -> Result<(), GovernanceError> {
         let old_delegate = self._delegates(&delegator);
 
-        self.data::<Data>().delegation.insert(&delegator, &delegatee);
+        self.data::<Data>()
+            .delegation
+            .insert(&delegator, &delegatee.ok_or(GovernanceError::InvalidInput)?);
 
-        self.emit_delegate_changed_event(&delegator, &old_delegate, &Some(delegatee.clone()));
+        self.emit_delegate_changed_event(&delegator, &old_delegate, delegatee);
 
         self._move_delegate_votes(
             &old_delegate,
-            &Some(delegatee.clone()),
-            self._get_voting_units(&delegator),
+            delegatee,
+            self._get_voting_units(&delegator.ok_or(GovernanceError::InvalidInput)?),
         )
     }
 
     fn _transfer_voting_units(
         &mut self,
-        from: &AccountId,
-        to: &AccountId,
+        from: &Option<AccountId>,
+        to: &Option<AccountId>,
         amount: Balance,
     ) -> Result<(), GovernanceError> {
         let mut store = self.data::<Data>().total_checkpoints.get_or_default();
-        if from == &AccountId::from([0x0; 32]) {
+        if from.is_none() {
             self._push(&mut store, Self::_add, amount)?;
         }
-        if to == &AccountId::from([0x0; 32]) {
+        if to.is_none() {
             self._push(&mut store, Self::_sub, amount)?;
         }
         self._move_delegate_votes(&self._delegates(from), &self._delegates(to), amount)

diff --git a/contracts/src/traits/governance/utils/votes.rs b/contracts/src/traits/governance/utils/votes.rs
@@ -37,4 +37,4 @@ pub trait Votes {
 }
 
 #[openbrush::wrapper]
-pub type VotesWrapper = dyn Votes;
+pub type VotesRef = dyn Votes;
diff --git a/examples/psp22_extensions/votes/lib.rs b/examples/psp22_extensions/votes/lib.rs
@@ -36,6 +36,7 @@ pub mod my_psp22_votes {
         votes: votes::Data,
         #[storage_field]
         nonces: nonces::Data,
+        mock_timestamp: Timestamp,
     }
 
     impl Contract {
@@ -45,8 +46,25 @@ pub mod my_psp22_votes {
 
             psp22::Internal::_mint_to(&mut instance, Self::env().caller(), total_supply).expect("Should mint");
 
+            instance.mock_timestamp = Self::env().block_timestamp();
+
             instance
         }
+
+        #[ink(message)]
+        pub fn block_timestamp(&self) -> Timestamp {
+            self.mock_timestamp
+        }
+
+        #[ink(message)]
+        pub fn set_block_timestamp(&mut self, timestamp: Timestamp) {
+            self.mock_timestamp = timestamp;
+        }
+
+        #[ink(message)]
+        pub fn increase_block_timestamp(&mut self, timestamp: Timestamp) {
+            self.mock_timestamp += timestamp;
+        }
     }
 
     impl NoncesImpl for Contract {}
@@ -59,7 +77,11 @@ pub mod my_psp22_votes {
         }
     }
 
-    impl TimestampProvider for Contract {}
+    impl TimestampProvider for Contract {
+        fn block_timestamp(&self) -> Timestamp {
+            self.mock_timestamp
+        }
+    }
 
     impl VotesImpl for Contract {}
 

diff --git a/tests/e2e/governance/governor.tests.ts b/tests/e2e/governance/governor.tests.ts
@@ -45,11 +45,13 @@ describe('Governor', function () {
     const contractAddressGovernance = (await contractFactoryGovernance.new(contractAddressVotes, votingDelay, votingPeriod, proposalThreshold, numrator)).address
     const contractGovernance = new ContractGovernance(contractAddressGovernance, deployer, api)
 
+    await contractVotes.tx.setBlockTimestamp((await contractGovernance.query.blockTimestamp()).value.ok!)
+
     const contractFactoryReceiver = new ConstructorsReceiver(api, deployer)
     const contractAddressReceiver = (await contractFactoryReceiver.new()).address
     const contractReceiver = new ContractReceiver(contractAddressReceiver, deployer, api)
 
-    const helper = new GovernorHelper(contractGovernance)
+    const helper = new GovernorHelper(contractGovernance, contractVotes)
 
     await helper.delegate(contractVotes, deployer, alice, 10)
     await helper.delegate(contractVotes, deployer, bob, 10)
@@ -212,7 +214,7 @@ describe('Governor', function () {
           '<description>#proposer=' + SS58ToHex(api, deployer.address)
         )
         await expect(helper.propose(deployer)).to.eventually.be.fulfilled
-        await helper.waitForDeadline()
+        await helper.waitForDeadline(1)
         await expect(helper.castVote(deployer, VoteType.for)).to.eventually.be.rejected
 
         await api.disconnect()
@@ -457,19 +459,10 @@ describe('Governor', function () {
     it('Succeeded', async function () {
       const {
         api,
-        bob,
         deployer,
-        contractVotes,
         helper
       } = await setup()
 
-      helper.addProposal(
-        contractVotes.address,
-        getSelectorByName(contractVotes.abi.messages, 'PSP22::transfer'),
-        [bob.address, new BN(1000), ''],
-        '<description>#proposer=' + SS58ToHex(api, deployer.address)
-      )
-
       await expect(helper.propose(deployer)).to.eventually.be.fulfilled
       await helper.waitForSnapshot()
 
@@ -487,18 +480,10 @@ describe('Governor', function () {
     it('Executed', async function () {
       const {
         api,
-        bob,
         deployer,
-        contractVotes,
         helper
       } = await setup()
 
-      helper.addProposal(
-        contractVotes.address,
-        getSelectorByName(contractVotes.abi.messages, 'PSP22::transfer'),
-        [bob.address, new BN(1000), ''],
-        '<description>#proposer=' + SS58ToHex(api, deployer.address)
-      )
       await expect(helper.propose(deployer)).to.eventually.be.fulfilled
       await helper.waitForSnapshot()
       await expect(helper.castVote(deployer, VoteType.for)).to.eventually.be.fulfilled
@@ -652,9 +637,12 @@ describe('Governor', function () {
         )
         await expect(helper.propose(deployer)).to.eventually.be.fulfilled
         await helper.waitForSnapshot()
+
         await expect(helper.castVote(deployer, VoteType.for)).to.eventually.be.fulfilled
         await helper.waitForDeadline(1)
+
         await expect(helper.execute(deployer)).to.eventually.be.fulfilled
+
         await expect(helper.cancel(deployer)).to.eventually.be.rejected
 
         await api.disconnect()

diff --git a/tests/e2e/governance/helper.ts b/tests/e2e/governance/helper.ts
@@ -12,10 +12,12 @@ export class GovernorHelper {
   private proposal: Transaction | undefined;
   private description: string | undefined;
   private governor: ContractGovernance | undefined;
+  private token: ContractVotes | undefined;
   private proposalId: number[] | undefined;
 
-  constructor(governor: ContractGovernance) {
+  constructor(governor: ContractGovernance, token: ContractVotes){
     this.governor = governor
+    this.token = token
   }
 
   addProposal(callee: string, selector: number[], input: (string | number | BN)[], description: string) {
@@ -60,9 +62,11 @@ export class GovernorHelper {
     }
 
     if(proposer) {
+      console.log((await this.governor?.query.propose([this.proposal!], this.description!))?.value)
       await this.governor?.withSigner(proposer).tx.propose([this.proposal!], this.description!)
     }
     else {
+      console.log((await this.governor?.query.propose([this.proposal!], this.description!))?.value)
       await this.governor?.tx.propose([this.proposal!], this.description!)
     }
   }
@@ -81,6 +85,7 @@ export class GovernorHelper {
     if(proposalSnapshot === undefined) throw new Error('Proposal snapshot not set')
 
     await this.governor?.tx.setBlockTimestamp(proposalSnapshot + offset)
+    await this.token?.tx.setBlockTimestamp(proposalSnapshot + offset)
   }
 
   async castVote(voter: KeyringPair, vote: VoteType) {
@@ -111,6 +116,7 @@ export class GovernorHelper {
     if(proposalDeadline === undefined) throw new Error('Proposal deadline not set')
 
     await this.governor?.tx.setBlockTimestamp(proposalDeadline + offset)
+    await this.token?.tx.setBlockTimestamp(proposalDeadline + offset)
   }
 
   async execute(proposer?: KeyringPair) {