diff --git a/src/interface.cairo b/src/interface.cairo index 862abb2..febf571 100644 --- a/src/interface.cairo +++ b/src/interface.cairo @@ -7,10 +7,10 @@ pub trait IPushComm { 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 index c02ace0..a89f542 100644 --- a/src/lib.cairo +++ b/src/lib.cairo @@ -42,10 +42,10 @@ pub mod PushComm { // Channels delegatedNotificationSenders: Map>, // 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 index 8763f89..892b3ec 100644 --- 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 index 3501fe0..bf23c02 100644 --- 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));