Safe upgrades: Add healthcheck function and call it when upgrading a contract #116 -Malachi PR

src/proposals.cairo

@@ -7,6 +7,8 @@ use starknet::ContractAddress;
 trait IProposals<TContractState> {
     fn vote(ref self: TContractState, prop_id: felt252, opinion: felt252);
     fn get_proposal_details(self: @TContractState, prop_id: felt252) -> PropDetails;
+    fn get_latest_proposal_id(self: @TContractState) -> felt252;
+
     fn get_vote_counts(self: @TContractState, prop_id: felt252) -> (u128, u128);
     fn submit_proposal(
         ref self: TContractState, payload: felt252, to_upgrade: ContractType
@@ -300,6 +302,10 @@ mod proposals {
             self.proposal_details.read(prop_id)
         }
 
+        fn get_latest_proposal_id(self: @ComponentState<TContractState>) -> felt252 {
+            self.get_free_prop_id_timestamp() - 1
+        }
+
         fn get_live_proposals(self: @ComponentState<TContractState>) -> Array<felt252> {
             let max: u32 = self.get_free_prop_id_timestamp().try_into().unwrap();
             let mut i: u32 = 0;

src/testing/setup.cairo

@@ -10,6 +10,7 @@ use konoha::contract::IGovernanceDispatcher;
 use konoha::contract::IGovernanceDispatcherTrait;
 use konoha::proposals::IProposalsDispatcher;
 use konoha::proposals::IProposalsDispatcherTrait;
+use konoha::staking::{IStakingDispatcher, IStakingDispatcherTrait};
 use konoha::upgrades::IUpgradesDispatcher;
 use konoha::upgrades::IUpgradesDispatcherTrait;
 use openzeppelin::token::erc20::interface::IERC20;
@@ -19,14 +20,17 @@ use snforge_std::{
 };
 use starknet::ContractAddress;
 use starknet::get_block_timestamp;
+//use super::staking_tests::set_floating_token_address;
 
-
-const GOV_TOKEN_INITIAL_SUPPLY: felt252 = 1000000000000000000;
+const GOV_TOKEN_INITIAL_SUPPLY: u256 = 1000000000000000000;
 
 const first_address: felt252 = 0x1;
 const second_address: felt252 = 0x2;
 const admin_addr: felt252 = 0x3;
 
+const governance_address: felt252 = 0x99999;
+
+// DEPRECATED, use deploy_governance_and_both_tokens instead
 fn deploy_governance(token_address: ContractAddress) -> IGovernanceDispatcher {
     let gov_contract = declare("Governance").expect('unable to declare governance');
     let mut args: Array<felt252> = ArrayTrait::new();
@@ -35,23 +39,49 @@ fn deploy_governance(token_address: ContractAddress) -> IGovernanceDispatcher {
     IGovernanceDispatcher { contract_address: address }
 }
 
+// return governance, voting token, floating token.
+// by default, all floating tokens are minted to admin address.
+fn deploy_governance_and_both_tokens() -> (
+    IGovernanceDispatcher, IERC20Dispatcher, IERC20Dispatcher
+) {
+    let gov_contract = declare("Governance").expect('unable to declare governance');
+    let floating_token_class = declare("FloatingToken").expect('unable to declare FloatingToken');
+    let voting_token_class = declare("VotingToken").expect('unable to declare VotingToken');
+    let mut args: Array<felt252> = ArrayTrait::new();
+    args.append(voting_token_class.class_hash.into());
+    args.append(floating_token_class.class_hash.into());
+    args.append(admin_addr);
+    gov_contract
+        .deploy_at(@args, governance_address.try_into().unwrap())
+        .expect('unable to deploy governance');
+    let gov_dispatcher = IGovernanceDispatcher {
+        contract_address: governance_address.try_into().unwrap()
+    };
+    let staking = IStakingDispatcher { contract_address: governance_address.try_into().unwrap() };
+
+    let voting_token_dispatcher = IERC20Dispatcher {
+        contract_address: gov_dispatcher.get_governance_token_address()
+    };
+
+    let floating_token_dispatcher = IERC20Dispatcher {
+        contract_address: staking.get_floating_token_address()
+    };
+
+    (gov_dispatcher, voting_token_dispatcher, floating_token_dispatcher)
+}
 
+// DEPRECATED, use deploy_governance_and_both_tokens instead
 fn deploy_and_distribute_gov_tokens(recipient: ContractAddress) -> IERC20Dispatcher {
     let mut calldata = ArrayTrait::new();
-    calldata.append(GOV_TOKEN_INITIAL_SUPPLY);
+    calldata.append(GOV_TOKEN_INITIAL_SUPPLY.low.into());
+    calldata.append(GOV_TOKEN_INITIAL_SUPPLY.high.into());
     calldata.append(recipient.into());
 
     let gov_token_contract = declare("FloatingToken").expect('unable to declare FloatingToken');
     let (token_addr, _) = gov_token_contract
         .deploy(@calldata)
         .expect('unable to deploy FloatingToken');
-    let token: IERC20Dispatcher = IERC20Dispatcher { contract_address: token_addr };
-
-    start_prank(CheatTarget::One(token_addr), admin_addr.try_into().unwrap());
-
-    token.transfer(first_address.try_into().unwrap(), 100000);
-    token.transfer(second_address.try_into().unwrap(), 100000);
-    token
+    IERC20Dispatcher { contract_address: token_addr }
 }
 
 
@@ -79,10 +109,29 @@ fn test_vote_upgrade_root(new_merkle_root: felt252) {
     assert(check_if_healthy(gov_contract_addr), 'new gov not healthy');
 }
 
-fn check_if_healthy(gov_contract_addr: ContractAddress) -> bool {
-    // TODO
-    let dispatcher = IProposalsDispatcher { contract_address: gov_contract_addr };
-    dispatcher.get_proposal_status(0);
-    let prop_details = dispatcher.get_proposal_details(0);
-    (prop_details.payload + prop_details.to_upgrade) != 0
+//if trying to apply_passed_proposal throws error:
+//Requested contract address ContractAddress(PatriciaKey(StarkFelt("0x0000000000000000000000000000000000000000000000000000000000000000"))) is not deployed.
+//if health check fails
+
+//health check completed for checking governance type. 
+//TODO: version history
+
+fn check_if_healthy(gov_address: ContractAddress) -> bool {
+    println!("Health contract address: {:?}", gov_address);
+
+    let proposals_dispatcher = IProposalsDispatcher { contract_address: gov_address };
+    let upgrades_dispatcher = IUpgradesDispatcher { contract_address: gov_address };
+
+    // Check if there are no proposals
+    let current_prop_id = proposals_dispatcher.get_latest_proposal_id();
+    if current_prop_id == 0 {
+        return true;
+    }
+    // Retrieve current proposal details
+    let current_prop_details = proposals_dispatcher.get_proposal_details(current_prop_id);
+
+    // Check if the latest upgrade type matches the proposal's upgrade type
+    let (_, last_upgrade_type) = upgrades_dispatcher.get_latest_upgrade();
+    last_upgrade_type.into() == current_prop_details.to_upgrade
 }
+

src/upgrades.cairo

@@ -1,6 +1,7 @@
 #[starknet::interface]
 trait IUpgrades<TContractState> {
     fn apply_passed_proposal(ref self: TContractState, prop_id: felt252);
+    fn get_latest_upgrade(self: @TContractState) -> (u64, u64);
 }
 
 #[starknet::component]
@@ -31,9 +32,11 @@ mod upgrades {
     #[storage]
     struct Storage {
         proposal_applied: LegacyMap::<felt252, bool>,
-        amm_address: ContractAddress
+        amm_address: ContractAddress,
+        latest_upgrade: (u64, u64), // (prop_id, upgrade_type)
     }
 
+
     #[derive(starknet::Event, Drop)]
     #[event]
     enum Event {
@@ -54,6 +57,9 @@ mod upgrades {
         impl Proposals: proposals_component::HasComponent<TContractState>,
         impl Airdrop: airdrop_component::HasComponent<TContractState>
     > of super::IUpgrades<ComponentState<TContractState>> {
+        fn get_latest_upgrade(self: @ComponentState<TContractState>) -> (u64, u64) {
+            self.latest_upgrade.read()
+        }
         fn apply_passed_proposal(ref self: ComponentState<TContractState>, prop_id: felt252) {
             let proposals_comp = get_dep_component!(@self, Proposals);
             let status = proposals_comp
@@ -136,6 +142,8 @@ mod upgrades {
                 _ => { panic_with_felt252('invalid to_upgrade') }
             };
             self.proposal_applied.write(prop_id, true); // Mark the proposal as applied
+            let upgrade_type: u64 = contract_type.try_into().unwrap();
+            self.latest_upgrade.write((prop_id.try_into().unwrap(), upgrade_type));
             self
                 .emit(
                     Upgraded {

tests/lib.cairo

@@ -1,10 +1,12 @@
-mod airdrop_tests;
-mod basic;
-mod proposals_tests;
+//mod airdrop_tests;
+//mod basic;
+//mod proposals_tests;
 mod setup;
 mod staking_tests;
-mod test_storage_pack;
-mod test_streaming;
-mod test_treasury;
-mod upgrades_tests;
-mod vesting;
+//mod test_storage_pack;
+//mod test_streaming;
+//mod test_treasury;
+//mod upgrades_tests;
+//mod vesting;
+mod test_setup;
+

tests/setup.cairo

@@ -20,7 +20,7 @@ use snforge_std::{
 };
 use starknet::ContractAddress;
 use starknet::get_block_timestamp;
-use super::staking_tests::set_floating_token_address;
+//use super::staking_tests::set_floating_token_address;
 
 const GOV_TOKEN_INITIAL_SUPPLY: u256 = 1000000000000000000;
 
@@ -109,10 +109,29 @@ fn test_vote_upgrade_root(new_merkle_root: felt252) {
     assert(check_if_healthy(gov_contract_addr), 'new gov not healthy');
 }
 
-fn check_if_healthy(gov_contract_addr: ContractAddress) -> bool {
-    // TODO
-    let dispatcher = IProposalsDispatcher { contract_address: gov_contract_addr };
-    dispatcher.get_proposal_status(0);
-    let prop_details = dispatcher.get_proposal_details(0);
-    (prop_details.payload + prop_details.to_upgrade) != 0
+//if trying to apply_passed_proposal throws error:
+//Requested contract address ContractAddress(PatriciaKey(StarkFelt("0x0000000000000000000000000000000000000000000000000000000000000000"))) is not deployed.
+//if health check fails
+
+//health check completed for checking governance type. 
+//TODO: version history
+
+fn check_if_healthy(gov_address: ContractAddress) -> bool {
+    println!("Health contract address: {:?}", gov_address);
+
+    let proposals_dispatcher = IProposalsDispatcher { contract_address: gov_address };
+    let upgrades_dispatcher = IUpgradesDispatcher { contract_address: gov_address };
+
+    // Check if there are no proposals
+    let current_prop_id = proposals_dispatcher.get_latest_proposal_id();
+    if current_prop_id == 0 {
+        return true;
+    }
+    // Retrieve current proposal details
+    let current_prop_details = proposals_dispatcher.get_proposal_details(current_prop_id);
+
+    // Check if the latest upgrade type matches the proposal's upgrade type
+    let (_, last_upgrade_type) = upgrades_dispatcher.get_latest_upgrade();
+    last_upgrade_type.into() == current_prop_details.to_upgrade
 }
+

tests/test_setup.cairo

@@ -0,0 +1,134 @@
+use array::ArrayTrait;
+use core::traits::Into;
+use core::traits::TryInto;
+use debug::PrintTrait;
+use konoha::constants;
+use konoha::contract::IGovernanceDispatcher;
+use konoha::contract::IGovernanceDispatcherTrait;
+use konoha::discussion::IDiscussionDispatcher;
+use konoha::discussion::IDiscussionDispatcherTrait;
+use konoha::proposals::IProposalsDispatcher;
+use konoha::proposals::IProposalsDispatcherTrait;
+use konoha::staking::{IStakingDispatcher, IStakingDispatcherTrait};
+use konoha::types::ContractType;
+use konoha::upgrades::{IUpgradesDispatcher, IUpgradesDispatcherTrait};
+use openzeppelin::token::erc20::interface::{
+    IERC20Dispatcher, IERC20DispatcherTrait, IERC20CamelOnlyDispatcher,
+    IERC20CamelOnlyDispatcherTrait
+};
+
+use snforge_std::{
+    BlockId, declare, ContractClassTrait, ContractClass, start_prank, start_warp, CheatTarget,
+    prank, CheatSpan, get_class_hash
+};
+use starknet::ContractAddress;
+
+use starknet::get_block_timestamp;
+
+use super::setup::check_if_healthy;
+
+use super::setup::{
+    admin_addr, first_address, second_address, deploy_governance, deploy_and_distribute_gov_tokens,
+    deploy_governance_and_both_tokens, test_vote_upgrade_root
+};
+use super::staking_tests::{set_staking_curve, stake_all, stake_half};
+
+#[test]
+fn test_healthy_upgrade() {
+    let (gov, _voting, floating) = deploy_governance_and_both_tokens();
+    set_staking_curve(gov.contract_address);
+    stake_all(gov.contract_address, floating, admin_addr.try_into().unwrap());
+
+    let proposals_dispatcher = IProposalsDispatcher { contract_address: gov.contract_address };
+    let upgrades_dispatcher = IUpgradesDispatcher { contract_address: gov.contract_address };
+
+    // Submit first proposal
+    start_prank(CheatTarget::One(gov.contract_address), admin_addr.try_into().unwrap());
+    let prop_id = proposals_dispatcher.submit_proposal(42, 3);
+
+    // Vote on the first proposal
+    start_prank(CheatTarget::One(gov.contract_address), admin_addr.try_into().unwrap());
+    proposals_dispatcher.vote(prop_id, 1);
+
+    // Check the status of the first proposal
+    assert_eq!(proposals_dispatcher.get_proposal_status(prop_id), 1, "Proposal not passed!");
+
+    let is_healthy = check_if_healthy(gov.contract_address);
+    assert!(is_healthy, "1Governance should be healthy bc empty");
+
+    upgrades_dispatcher.apply_passed_proposal(prop_id);
+
+    // Submit second proposal
+    start_prank(CheatTarget::One(gov.contract_address), admin_addr.try_into().unwrap());
+    let prop_id1 = proposals_dispatcher.submit_proposal(43, 3);
+
+    // Vote on the second proposal
+    start_prank(CheatTarget::One(gov.contract_address), admin_addr.try_into().unwrap());
+    proposals_dispatcher.vote(prop_id1, 1);
+
+    // Check the status of the second proposal
+    assert_eq!(
+        proposals_dispatcher.get_proposal_status(prop_id1), 1, "second Proposal not passed!"
+    );
+
+    let is_healthy_after = check_if_healthy(gov.contract_address);
+    assert!(
+        is_healthy_after, "2-Governance should be healthy after same type to type (3) upgrade."
+    );
+
+    upgrades_dispatcher.apply_passed_proposal(prop_id1);
+}
+#[test]
+fn test_unhealthy_upgrade() {
+    // Deploy governance and tokens
+    let (gov, _voting, floating) = deploy_governance_and_both_tokens();
+    let gov_address = gov.contract_address;
+
+    // Print the governance address for debugging
+    println!("Governance contract address: {:?}", gov_address);
+
+    // Initialize staking curve
+    set_staking_curve(gov_address);
+
+    // Stake tokens
+    stake_all(gov_address, floating, admin_addr.try_into().unwrap());
+
+    // Dispatcher for proposals
+    let dispatcher = IProposalsDispatcher { contract_address: gov_address };
+
+    // Submit first proposal
+    start_prank(CheatTarget::One(gov_address), admin_addr.try_into().unwrap());
+    let prop_id = dispatcher.submit_proposal(42, 3);
+    dispatcher.vote(prop_id, 1);
+    assert_eq!(dispatcher.get_proposal_status(prop_id), 1, "First proposal not passed!");
+
+    // Check health (should be healthy)
+    let is_healthy = check_if_healthy(gov_address);
+    println!("After first proposal, is_healthy: {}", is_healthy);
+    assert!(is_healthy, "Governance should be healthy after first proposal");
+
+    // Apply the first proposal
+    IUpgradesDispatcher { contract_address: gov_address }.apply_passed_proposal(prop_id);
+
+    // Submit second proposal (different type)
+    start_prank(CheatTarget::One(gov_address), admin_addr.try_into().unwrap());
+    let prop_id2 = dispatcher.submit_proposal(43, 5);
+    dispatcher.vote(prop_id2, 1);
+    assert_eq!(dispatcher.get_proposal_status(prop_id2), 1, "Second proposal not passed!");
+
+    // Print details for debugging
+    let proposals_dispatcher = IProposalsDispatcher { contract_address: gov_address };
+    let upgrades_dispatcher = IUpgradesDispatcher { contract_address: gov_address };
+
+    let (_, last_upgrade_type) = upgrades_dispatcher.get_latest_upgrade();
+    let current_prop_id = proposals_dispatcher.get_latest_proposal_id();
+    let current_prop_details = proposals_dispatcher.get_proposal_details(current_prop_id);
+
+    println!("Governance Type: {:?}", last_upgrade_type);
+    println!("Upgrading Type: {:?}", current_prop_details.to_upgrade);
+
+    // Check health after the second proposal
+    let is_healthy_after = check_if_healthy(gov_address);
+    println!("After second proposal, is_healthy: {}", is_healthy_after);
+    assert!(!is_healthy_after, "Governance should not be healthy after the second proposal");
+}