feat: governance & push token address

push-protocol · Aug 18, 2024 · 3c414b8 · 3c414b8
1 parent 833b8f2
commit 3c414b8
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 16 deletions.
diff --git a/src/interface.cairo b/src/interface.cairo
@@ -7,10 +7,10 @@ pub trait IPushComm<TContractState> {
     fn get_migration_status(self: @TContractState) -> bool;
     fn set_push_core_address(ref self: TContractState, core_address: felt252);
     fn get_push_core_address(self: @TContractState) -> felt252;
-    fn get_push_governance_address(self: @TContractState) -> felt252;
-    fn set_push_governance_address(ref self: TContractState, governance_address: felt252);
-    fn get_push_token_address(self: @TContractState) -> felt252;
-    fn set_push_token_address(ref self: TContractState, push_token_address: felt252);
+    fn get_push_governance_address(self: @TContractState) -> ContractAddress;
+    fn set_push_governance_address(ref self: TContractState, governance_address: ContractAddress);
+    fn get_push_token_address(self: @TContractState) -> ContractAddress;
+    fn set_push_token_address(ref self: TContractState, push_token_address: ContractAddress);
     // Channel
     fn verify_channel_alias(ref self: TContractState, channel_address: felt252);
     fn add_delegate(ref self: TContractState, delegate: ContractAddress);

diff --git a/src/lib.cairo b/src/lib.cairo
@@ -42,10 +42,10 @@ pub mod PushComm {
         // Channels
         delegatedNotificationSenders: Map<ContractAddress, Map<ContractAddress, bool>>,
         // Contract State
-        governance: felt252,
+        governance: ContractAddress,
         is_migration_complete: bool,
         push_core_address: felt252,
-        push_token_address: felt252,
+        push_token_address: ContractAddress,
         // Chain Info
         chain_name: felt252,
         chain_id: felt252,
@@ -132,12 +132,18 @@ pub mod PushComm {
 
 
     #[constructor]
-    fn constructor(ref self: ContractState, owner: ContractAddress, chain_name: felt252) {
-        // Set the initial owner of the contract
-        self.ownable.initializer(owner);
+    fn constructor(
+        ref self: ContractState,
+        owner: ContractAddress,
+        push_governance: ContractAddress,
+        chain_name: felt252
+    ) {
         let chain_id = get_execution_info().unbox().tx_info.unbox().chain_id;
+
+        self.ownable.initializer(owner);
         self.chain_id.write(chain_id);
         self.chain_name.write(chain_name);
+        self.governance.write(push_governance);
     }
 
     #[generate_trait]
@@ -343,21 +349,23 @@ pub mod PushComm {
 
 
         // Infos
-        fn set_push_governance_address(ref self: ContractState, governance_address: felt252) {
+        fn set_push_governance_address(
+            ref self: ContractState, governance_address: ContractAddress
+        ) {
             self.ownable.assert_only_owner();
             self.governance.write(governance_address);
         }
 
-        fn get_push_governance_address(self: @ContractState) -> felt252 {
+        fn get_push_governance_address(self: @ContractState) -> ContractAddress {
             self.governance.read()
         }
 
-        fn set_push_token_address(ref self: ContractState, push_token_address: felt252) {
+        fn set_push_token_address(ref self: ContractState, push_token_address: ContractAddress) {
             self.ownable.assert_only_owner();
             self.push_token_address.write(push_token_address);
         }
 
-        fn get_push_token_address(self: @ContractState) -> felt252 {
+        fn get_push_token_address(self: @ContractState) -> ContractAddress {
             self.push_token_address.read()
         }
     }

diff --git a/tests/common.cairo b/tests/common.cairo
@@ -6,6 +6,11 @@ pub fn PUSH_ADMIN() -> ContractAddress {
     'push_admin'.try_into().unwrap()
 }
 
+pub fn PUSH_GOVERNANCE() -> ContractAddress {
+    'push_governance'.try_into().unwrap()
+}
+
+
 pub fn USER_1() -> ContractAddress {
     'user_1'.try_into().unwrap()
 }
@@ -17,7 +22,7 @@ pub fn CHAIN_NAME() -> felt252 {
 pub fn deploy_contract() -> ContractAddress {
     let contract = declare("PushComm").unwrap();
 
-    let calldata = array![PUSH_ADMIN().into(), CHAIN_NAME()];
+    let calldata = array![PUSH_ADMIN().into(), PUSH_GOVERNANCE().into(), CHAIN_NAME()];
     let (contract_address, _) = contract.deploy(@calldata).unwrap();
     contract_address
 }
diff --git a/tests/test_push_admin.cairo b/tests/test_push_admin.cairo
@@ -42,7 +42,7 @@ fn test_gov_address() {
     let contract_address = deploy_contract();
     let push_comm = IPushCommDispatcher { contract_address };
 
-    let GOV_ADDRESS: felt252 = 'some addrs';
+    let GOV_ADDRESS: ContractAddress = 'some addrs'.try_into().unwrap();
 
     // admin sets the migration status
     cheat_caller_address(contract_address, PUSH_ADMIN(), CheatSpan::TargetCalls(1));
@@ -57,7 +57,7 @@ fn test_push_token_address() {
     let contract_address = deploy_contract();
     let push_comm = IPushCommDispatcher { contract_address };
 
-    let TOKEN_ADDRESS: felt252 = 'some addrs';
+    let TOKEN_ADDRESS: ContractAddress = 'user_1'.try_into().unwrap();
 
     // admin sets the migration status
     cheat_caller_address(contract_address, PUSH_ADMIN(), CheatSpan::TargetCalls(1));