add method to check how a particular account voted (#51)

EkuboProtocol · May 28, 2024 · cb8c076 · cb8c076
1 parent dfd0a69
commit cb8c076
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 4 deletions.
diff --git a/src/governor.cairo b/src/governor.cairo
@@ -86,6 +86,12 @@ pub trait IGovernor<TContractState> {
     // Gets the proposal and the config version with which it was created
     fn get_proposal_with_config(self: @TContractState, id: felt252) -> (ProposalInfo, Config);
 
+    // Returns the vote cast by the given voter on the given proposal ID.
+    // - 0 means not voted, or proposal does not exist
+    // - 3 means voted in favor
+    // - 1 means voted against
+    fn get_vote(self: @TContractState, id: felt252, voter: ContractAddress) -> u8;
+
     // Change the configuration of the governor. Only affects proposals created after the configuration change. Must be called by self, e.g. via a proposal.
     fn reconfigure(ref self: TContractState, config: Config) -> u64;
 
@@ -173,7 +179,7 @@ pub mod Governor {
         config: Config,
         nonce: u64,
         proposals: LegacyMap<felt252, ProposalInfo>,
-        has_voted: LegacyMap<(ContractAddress, felt252), bool>,
+        vote: LegacyMap<(felt252, ContractAddress), u8>,
         latest_proposal_by_proposer: LegacyMap<ContractAddress, felt252>,
         latest_config_version: u64,
         config_versions: LegacyMap<u64, Config>,
@@ -295,11 +301,11 @@ pub mod Governor {
             let timestamp_current = get_block_timestamp();
             let voting_start_time = (proposal.execution_state.created + config.voting_start_delay);
             let voter = get_caller_address();
-            let has_voted = self.has_voted.read((voter, id));
+            let past_vote = self.vote.read((id, voter));
 
             assert(timestamp_current >= voting_start_time, 'VOTING_NOT_STARTED');
             assert(timestamp_current < (voting_start_time + config.voting_period), 'VOTING_ENDED');
-            assert(!has_voted, 'ALREADY_VOTED');
+            assert(past_vote.is_zero(), 'ALREADY_VOTED');
 
             let weight = self
                 .get_staker()
@@ -315,7 +321,11 @@ pub mod Governor {
                 proposal.nay = proposal.nay + weight;
             }
             self.proposals.write(id, proposal);
-            self.has_voted.write((voter, id), true);
+            self.vote.write((id, voter), if yea {
+                3
+            } else {
+                1
+            });
 
             self.emit(Voted { id, voter, weight, yea });
         }
@@ -477,6 +487,10 @@ pub mod Governor {
             (proposal, config)
         }
 
+        fn get_vote(self: @ContractState, id: felt252, voter: ContractAddress) -> u8 {
+            self.vote.read((id, voter))
+        }
+
         fn reconfigure(ref self: ContractState, config: Config) -> u64 {
             self.check_self_call();
 

diff --git a/src/governor_test.cairo b/src/governor_test.cairo
@@ -450,6 +450,7 @@ fn test_vote_already_voted_should_fail() {
 
     set_contract_address(voter1());
     governor.vote(id, true);
+    assert_eq!(governor.get_vote(id, voter1()), 3);
 
     // Trying to vote twice on the same proposal should fail
     governor.vote(id, true);
@@ -738,6 +739,7 @@ fn test_execute_no_majority_should_fail() {
 
     set_contract_address(voter2());
     governor.vote(id, false);
+    assert_eq!(governor.get_vote(id, voter2()), 1);
     let proposal = governor.get_proposal(id);
     assert_eq!(proposal.yea, config.quorum);
     assert_eq!(proposal.nay, config.quorum + 1);