Add vesting tests (#43)

* Add vesting to the main contract

* Add tests

* Polish scarb fmt

* Update Scarb in CI too to fix CI

* Implement PR feedback

* Fix vesting tests with snforge v0.23

Scarb.toml

@@ -17,6 +17,9 @@ snforge_std = { git = "https://github.com/foundry-rs/starknet-foundry.git", tag
 
 [[target.starknet-contract]]
 
+[scripts]
+test = "snforge test"
+
 [[tool.snforge.fork]]
 name = "MAINNET"
 url = "http://34.22.208.73:6060/v0_7"

src/contract.cairo

@@ -18,6 +18,9 @@ trait IGovernance<TContractState> {
 
 // in component
 
+// VESTING
+
+// in component
 // OPTIONS / ONE-OFF
 }
 
@@ -31,17 +34,21 @@ mod Governance {
     use konoha::proposals::proposals as proposals_component;
     use konoha::upgrades::upgrades as upgrades_component;
     use konoha::airdrop::airdrop as airdrop_component;
+    use konoha::vesting::vesting as vesting_component;
 
     use starknet::ContractAddress;
 
 
     component!(path: airdrop_component, storage: airdrop, event: AirdropEvent);
+    component!(path: vesting_component, storage: vesting, event: VestingEvent);
     component!(path: proposals_component, storage: proposals, event: ProposalsEvent);
     component!(path: upgrades_component, storage: upgrades, event: UpgradesEvent);
 
     #[abi(embed_v0)]
     impl Airdrop = airdrop_component::AirdropImpl<ContractState>;
 
+    #[abi(embed_v0)]
+    impl Vesting = vesting_component::VestingImpl<ContractState>;
     #[abi(embed_v0)]
     impl Proposals = proposals_component::ProposalsImpl<ContractState>;
 
@@ -53,10 +60,12 @@ mod Governance {
         proposal_initializer_run: LegacyMap::<u64, bool>,
         governance_token_address: ContractAddress,
         #[substorage(v0)]
-        proposals: proposals_component::Storage,
-        #[substorage(v0)]
         airdrop: airdrop_component::Storage,
         #[substorage(v0)]
+        vesting: vesting_component::Storage,
+        #[substorage(v0)]
+        proposals: proposals_component::Storage,
+        #[substorage(v0)]
         upgrades: upgrades_component::Storage
     }
 
@@ -66,7 +75,7 @@ mod Governance {
     struct Proposed {
         prop_id: felt252,
         payload: felt252,
-        to_upgrade: ContractType
+        to_upgrade: ContractType,
     }
 
     #[derive(starknet::Event, Drop)]
@@ -82,6 +91,7 @@ mod Governance {
         Proposed: Proposed,
         Voted: Voted,
         AirdropEvent: airdrop_component::Event,
+        VestingEvent: vesting_component::Event,
         ProposalsEvent: proposals_component::Event,
         UpgradesEvent: upgrades_component::Event
     }

src/govtoken.cairo

@@ -0,0 +1,102 @@
+use starknet::{ContractAddress, ClassHash};
+
+#[starknet::interface]
+trait IveCARM<TContractState> {
+    fn name(self: @TContractState) -> felt252;
+    fn symbol(self: @TContractState) -> felt252;
+    fn decimals(self: @TContractState) -> u8;
+    fn mint(ref self: TContractState, recipient: ContractAddress, amount: u256);
+    fn burn(ref self: TContractState, account: ContractAddress, amount: u256);
+    fn upgrade(ref self: TContractState, new_class_hash: ClassHash);
+}
+
+#[starknet::contract]
+mod MyToken {
+    use openzeppelin::access::ownable::ownable::OwnableComponent::InternalTrait;
+    use starknet::ContractAddress;
+    use starknet::ClassHash;
+    use openzeppelin::token::erc20::ERC20Component;
+    use openzeppelin::access::ownable::ownable::OwnableComponent;
+
+    component!(path: ERC20Component, storage: erc20, event: ERC20Event);
+    component!(path: OwnableComponent, storage: ownable, event: OwnableEvent);
+
+    // ERC20 Component
+    #[abi(embed_v0)]
+    impl ERC20Impl = ERC20Component::ERC20Impl<ContractState>;
+
+    #[abi(embed_v0)]
+    impl ERC20CamelOnlyImpl = ERC20Component::ERC20CamelOnlyImpl<ContractState>;
+
+    impl ERC20InternalImpl = ERC20Component::InternalImpl<ContractState>;
+
+    impl OwnableInternalImpl = OwnableComponent::InternalImpl<ContractState>;
+
+    // Ownable Component
+    #[abi(embed_v0)]
+    impl OwnableImpl = OwnableComponent::OwnableImpl<ContractState>;
+
+
+    #[storage]
+    struct Storage {
+        #[substorage(v0)]
+        erc20: ERC20Component::Storage,
+        #[substorage(v0)]
+        ownable: OwnableComponent::Storage,
+    }
+
+    #[event]
+    #[derive(Drop, starknet::Event)]
+    enum Event {
+        Upgraded: Upgraded,
+        // #[flat]
+        ERC20Event: ERC20Component::Event,
+        OwnableEvent: OwnableComponent::Event
+    }
+
+    #[derive(Drop, starknet::Event)]
+    struct Upgraded {
+        class_hash: ClassHash
+    }
+
+
+    #[constructor]
+    fn constructor(ref self: ContractState, owner: ContractAddress) {
+        self.ownable.initializer(owner);
+    }
+
+    #[abi(embed_v0)]
+    impl VeCARMImpl of super::IveCARM<ContractState> {
+        // Did not import Erc20MetaData, so we can change decimals
+        // so we need to define name, symbol and decimals ourselves
+        fn name(self: @ContractState) -> felt252 {
+            'vote escrowed Carmine Token'
+        }
+
+        fn symbol(self: @ContractState) -> felt252 {
+            'veCARM'
+        }
+
+        fn decimals(self: @ContractState) -> u8 {
+            18
+        }
+
+        fn mint(ref self: ContractState, recipient: ContractAddress, amount: u256) {
+            self.ownable.assert_only_owner();
+            self.erc20._mint(recipient, amount);
+        }
+
+        fn burn(ref self: ContractState, account: ContractAddress, amount: u256) {
+            self.ownable.assert_only_owner();
+            self.erc20._burn(account, amount);
+        }
+
+        fn upgrade(ref self: ContractState, new_class_hash: ClassHash) {
+            self.ownable.assert_only_owner();
+            assert(!new_class_hash.is_zero(), 'Class hash cannot be zero');
+            starknet::replace_class_syscall(new_class_hash).unwrap();
+            self.emit(Upgraded { class_hash: new_class_hash });
+        }
+    }
+}
+

src/lib.cairo

@@ -12,6 +12,8 @@ mod traits;
 mod treasury;
 mod types;
 mod upgrades;
+mod vesting;
+mod govtoken; // if I put this in tests/ , I seem unable to use declare('MyToken')
 mod voting_token;
 mod testing {
     mod setup;

src/vesting.cairo

@@ -15,19 +15,21 @@ trait IVesting<TContractState> {
         ref self: TContractState,
         first_vest: u64,
         period: u64,
-        increments_count: u64,
-        total_amount: u128
+        increments_count: u16,
+        total_amount: u128,
+        grantee: ContractAddress
     );
 // MAYBE – streaming?
 // MAYBE – options on the govtoken?
 }
 
 #[starknet::component]
 mod vesting {
-    use starknet::syscalls::get_block_timestamp;
-
-    use konoha::traits::IGovernanceTokenDispatcher;
-    use konoha::traits::IGovernanceTokenDispatcherTrait;
+    use starknet::ContractAddress;
+    use starknet::{get_block_timestamp, get_caller_address, get_contract_address};
+    use konoha::contract::Governance;
+    use konoha::contract::{IGovernanceDispatcher, IGovernanceDispatcherTrait};
+    use konoha::traits::{IGovernanceTokenDispatcher, IGovernanceTokenDispatcherTrait};
 
     #[storage]
     struct Storage {
@@ -65,10 +67,11 @@ mod vesting {
             vested_timestamp: u64
         ) {
             let amt_to_vest = self.milestone.read((vested_timestamp, grantee));
-            assert(amt_to_vest != 0, 'no vesting milestone found, or already vested');
+            assert(amt_to_vest != 0, 'nothing to vest');
             assert(get_block_timestamp() > vested_timestamp, 'not yet eligible');
-            IGovernanceTokenDispatcher { contract_address: govtoken_addr }
-                .mint(claimee, u256 { high: 0, low: amt_to_vest });
+            let self_dsp = IGovernanceDispatcher { contract_address: get_contract_address() };
+            IGovernanceTokenDispatcher { contract_address: self_dsp.get_governance_token_address() }
+                .mint(grantee, amt_to_vest.into());
             self.milestone.write((vested_timestamp, grantee), 0);
             self
                 .emit(
@@ -82,29 +85,30 @@ mod vesting {
             grantee: ContractAddress,
             amount: u128
         ) {
-            self.milestone.write((vested_timestamp, grantee), amount);
+            assert(get_caller_address() == get_contract_address(), 'not self-call');
+            self.milestone.write((vesting_timestamp, grantee), amount);
             self
                 .emit(
                     VestingMilestoneAdded {
-                        grantee: grantee, timestamp: vesting_timestamp, amount: u128
+                        grantee: grantee, timestamp: vesting_timestamp, amount: amount
                     }
                 )
         }
 
         fn add_linear_vesting_schedule(
-            ref self: TContractState,
+            ref self: ComponentState<TContractState>,
             first_vest: u64,
             period: u64,
             increments_count: u16,
             total_amount: u128,
             grantee: ContractAddress
         ) {
+            assert(get_caller_address() == get_contract_address(), 'not self-call');
             let mut i: u16 = 0;
             let mut curr_timestamp = first_vest;
-            assert(increments_count > 1, 'schedule must have more than one milestone');
-            assert(get_block_timestamp() < first_vest, 'first vest cannot be in the past');
-            assert()
-            let per_vest_amount = total_amount / increments_count;
+            assert(increments_count > 1, 'increments_count <= 1');
+            assert(get_block_timestamp() < first_vest, 'first vest can\'t be in the past');
+            let per_vest_amount = total_amount / increments_count.into();
             let mut total_scheduled = 0;
             loop {
                 if i == increments_count {

tests/lib.cairo

@@ -1,2 +1,3 @@
+mod vesting;
 mod basic;
 mod test_treasury;

tests/vesting.cairo

@@ -0,0 +1,114 @@
+use core::option::OptionTrait;
+use core::result::ResultTrait;
+use array::ArrayTrait;
+use core::traits::TryInto;
+use debug::PrintTrait;
+use starknet::ContractAddress;
+use snforge_std::{
+    BlockId, declare, ContractClassTrait, ContractClass, start_prank, start_warp, CheatTarget
+};
+
+use konoha::vesting::{IVestingDispatcher, IVestingDispatcherTrait, IVesting};
+use openzeppelin::token::erc20::interface::{IERC20Dispatcher, IERC20DispatcherTrait};
+
+// returns gov addr, token addr
+fn test_setup() -> (ContractAddress, ContractAddress) {
+    let new_gov_contract: ContractClass = declare("Governance")
+        .expect('unable to declare Governance');
+    let new_token_contract: ContractClass = declare("MyToken").expect('unable to declare MyToken');
+    let new_gov_addr: ContractAddress =
+        0x001405ab78ab6ec90fba09e6116f373cda53b0ba557789a4578d8c1ec374ba0f
+        .try_into()
+        .unwrap();
+    let mut token_constructor = ArrayTrait::new();
+    token_constructor.append(new_gov_addr.into()); // Owner
+    let (token_address, _) = new_token_contract
+        .deploy(@token_constructor)
+        .expect('unable to deploy token');
+    let mut gov_constructor: Array<felt252> = ArrayTrait::new();
+    gov_constructor.append(token_address.into());
+    let (gov_address, _) = new_gov_contract
+        .deploy_at(@gov_constructor, new_gov_addr)
+        .expect('unable to deploy gov');
+
+    (gov_address, token_address)
+}
+
+#[test]
+#[should_panic(expected: ('not self-call',))]
+fn test_unauthorized_add_vesting_schedule() {
+    let (gov_address, _) = test_setup();
+
+    let gov_vesting = IVestingDispatcher { contract_address: gov_address };
+
+    start_warp(CheatTarget::All, 1);
+
+    gov_vesting.add_linear_vesting_schedule(10, 10, 10, 1000000, 0x1.try_into().unwrap());
+}
+
+#[test]
+#[should_panic(expected: ('not yet eligible',))]
+fn test_unauthorized_vest_early() {
+    let (gov_address, _) = test_setup();
+
+    let gov_vesting = IVestingDispatcher { contract_address: gov_address };
+
+    start_warp(CheatTarget::All, 1);
+    start_prank(CheatTarget::One(gov_address), gov_address);
+
+    let grantee: ContractAddress = 0x1.try_into().unwrap();
+
+    gov_vesting.add_linear_vesting_schedule(10, 10, 10, 1000000, grantee);
+
+    gov_vesting.vest(grantee, 10);
+}
+
+#[test]
+#[should_panic(expected: ('nothing to vest',))]
+fn test_vest_twice() {
+    let (gov_address, _) = test_setup();
+
+    let gov_vesting = IVestingDispatcher { contract_address: gov_address };
+
+    start_warp(CheatTarget::All, 1);
+    start_prank(CheatTarget::One(gov_address), gov_address);
+
+    let grantee: ContractAddress = 0x1.try_into().unwrap();
+
+    gov_vesting.add_linear_vesting_schedule(10, 10, 10, 1000000, grantee);
+
+    start_warp(CheatTarget::All, 11);
+
+    gov_vesting.vest(grantee, 10);
+    gov_vesting.vest(grantee, 10);
+}
+
+#[test]
+fn test_add_simple_vesting_schedule() {
+    let (gov_address, token_address) = test_setup();
+
+    let gov_vesting = IVestingDispatcher { contract_address: gov_address };
+    let tok = IERC20Dispatcher { contract_address: token_address };
+
+    start_warp(CheatTarget::All, 1);
+    start_prank(CheatTarget::One(gov_address), gov_address);
+
+    let grantee: ContractAddress = 0x1.try_into().unwrap();
+    gov_vesting.add_linear_vesting_schedule(10, 10, 10, 1000001, grantee);
+
+    start_warp(CheatTarget::All, 11); // past first vest
+    // anyone can claim for the grantee
+    gov_vesting.vest(grantee, 10);
+    assert(tok.balance_of(grantee) == 100000, 'vesting unsuccessful');
+
+    // grantee themselves can claim too
+    start_prank(CheatTarget::One(gov_address), grantee);
+    start_warp(CheatTarget::All, 21); // past second vest
+    gov_vesting.vest(grantee, 20);
+    assert(tok.balance_of(grantee) == 200000, 'vesting unsuccessful');
+
+    start_warp(CheatTarget::All, 101); // past last vest. no requirement to vest in order
+    gov_vesting.vest(grantee, 100);
+    // leftover tokens are included in last vest. (remainder after division)
+    assert(tok.balance_of(grantee) == 300001, 'vesting unsuccessful');
+}