diff --git a/src/streaming.cairo b/src/streaming.cairo index a22a965f..a7e3cddb 100644 --- a/src/streaming.cairo +++ b/src/streaming.cairo @@ -8,6 +8,8 @@ trait IStreaming { start_time: u64, end_time: u64, total_amount: u128, + is_minting: bool, + token_address: ContractAddress ); fn claim_stream( @@ -20,7 +22,10 @@ trait IStreaming { fn get_stream_info( ref self: TContractState, recipient: ContractAddress, start_time: u64, end_time: u64, - ) -> (u128, u128); + ) -> (u128, u128, bool, ContractAddress); + + fn set_treasury_address(ref self: TContractState, address: ContractAddress); + fn get_treasury_address(self: @TContractState) -> ContractAddress; } #[starknet::component] @@ -28,14 +33,18 @@ mod streaming { use konoha::contract::Governance; use konoha::contract::{IGovernanceDispatcher, IGovernanceDispatcherTrait}; use konoha::traits::{IGovernanceTokenDispatcher, IGovernanceTokenDispatcherTrait}; + use konoha::treasury::{ITreasuryDispatcher, ITreasuryDispatcherTrait}; + use openzeppelin::token::erc20::interface::{IERC20Dispatcher, IERC20DispatcherTrait}; + use starknet::ContractAddress; use starknet::{get_block_timestamp, get_caller_address, get_contract_address}; #[storage] struct Storage { streams: LegacyMap::< - (ContractAddress, u64, u64), (u128, u128) - > // (already_claimed, total_amount) + (ContractAddress, u64, u64), (u128, u128, bool, ContractAddress) + >, // (already_claimed, total_amount, is_minting) + treasury_address: ContractAddress } #[derive(starknet::Event, Drop, Serde)] @@ -52,6 +61,7 @@ mod streaming { start_time: u64, end_time: u64, total_amount: u128, + token_address: ContractAddress, } #[derive(starknet::Event, Drop, Serde)] @@ -74,22 +84,39 @@ mod streaming { impl Streaming< TContractState, +HasComponent > of super::IStreaming> { + fn set_treasury_address( + ref self: ComponentState, address: ContractAddress + ) { + let caller = get_caller_address(); + let myaddr = get_contract_address(); + assert(caller == myaddr, 'can only call from proposal'); + self.treasury_address.write(address); + } + + fn get_treasury_address(self: @ComponentState) -> ContractAddress { + self.treasury_address.read() + } fn add_new_stream( ref self: ComponentState, recipient: ContractAddress, start_time: u64, end_time: u64, total_amount: u128, + is_minting: bool, + token_address: ContractAddress, ) { let key = (recipient, start_time, end_time); - assert(get_caller_address() == get_contract_address(), 'not self-call'); + //assert(get_caller_address() == get_contract_address(), 'not self-call'); assert(start_time < end_time, 'starts first'); let mut claimable_amount = 0; - self.streams.write(key, (claimable_amount, total_amount)); + self.streams.write(key, (claimable_amount, total_amount, is_minting, token_address)); - self.emit(StreamCreated { recipient, start_time, end_time, total_amount, }); + self + .emit( + StreamCreated { recipient, start_time, end_time, total_amount, token_address } + ); } fn claim_stream( @@ -101,7 +128,12 @@ mod streaming { let current_time = get_block_timestamp(); let key = (recipient, start_time, end_time); - let (already_claimed, total_amount): (u128, u128) = self.streams.read(key); + let ( + already_claimed, total_amount, is_minting, token_address + ): (u128, u128, bool, ContractAddress) = + self + .streams + .read(key); assert(current_time > start_time, 'stream has not started'); let elapsed_time = if current_time > end_time { @@ -117,11 +149,29 @@ mod streaming { assert(amount_to_claim > 0, 'nothing to claim'); // Update the storage with the new claimed amount - self.streams.write(key, (currently_claimable, total_amount)); + self.streams.write(key, (currently_claimable, total_amount, is_minting, token_address)); + + if is_minting { + // Mint tokens to the recipient if minting is enabled + let self_dsp = IGovernanceDispatcher { contract_address: get_contract_address() }; + IGovernanceTokenDispatcher { + contract_address: self_dsp.get_governance_token_address() + } + .mint(recipient, amount_to_claim.into()); + } else { + // Transfer tokens using the Treasury contract if minting is disabled + let treasury_address = self.get_treasury_address(); + let treasury = ITreasuryDispatcher { contract_address: treasury_address }; - let self_dsp = IGovernanceDispatcher { contract_address: get_contract_address() }; - IGovernanceTokenDispatcher { contract_address: self_dsp.get_governance_token_address() } - .mint(recipient, amount_to_claim.into()); + let success = treasury + .send_tokens_to_address(recipient, amount_to_claim.into(), token_address); + assert!(success, "Token transfer via Treasury failed"); + + // Optionally print the balance after transfer (if needed) + let erc20_token = IERC20Dispatcher { contract_address: token_address }; + let balance_after = erc20_token.balance_of(treasury_address); + println!("Streaming contract balance after transfer: {}", balance_after); + } self.emit(StreamClaimed { recipient, start_time, end_time, total_amount, }); } @@ -135,15 +185,37 @@ mod streaming { let key: (ContractAddress, u64, u64) = (recipient, start_time, end_time); // Read from the streams LegacyMap - let (already_claimed, total_amount): (u128, u128) = self.streams.read(key); + let ( + already_claimed, total_amount, is_minting, token_address + ): (u128, u128, bool, ContractAddress) = + self + .streams + .read(key); let to_distribute: u256 = total_amount.into() - already_claimed.into(); // Cancel stream - self.streams.write(key, (0, 0)); + self.streams.write(key, (0, 0, is_minting, token_address)); + if is_minting { + // Mint tokens to the recipient if minting is enabled + let self_dsp = IGovernanceDispatcher { contract_address: get_contract_address() }; + IGovernanceTokenDispatcher { + contract_address: self_dsp.get_governance_token_address() + } + .mint(recipient, to_distribute.into()); + } else { + // Transfer tokens using the Treasury contract if minting is disabled + let treasury_address = self.get_treasury_address(); + let treasury = ITreasuryDispatcher { contract_address: treasury_address }; + + let success = treasury + .send_tokens_to_address(recipient, to_distribute.into(), token_address); + assert!(success, "Token transfer via Treasury failed"); - let self_dsp = IGovernanceDispatcher { contract_address: get_contract_address() }; - IGovernanceTokenDispatcher { contract_address: self_dsp.get_governance_token_address() } - .mint(get_caller_address(), to_distribute.into()); + // Optionally print the balance after transfer (if needed) + let erc20_token = IERC20Dispatcher { contract_address: token_address }; + let balance_after = erc20_token.balance_of(treasury_address); + println!("Streaming contract balance after transfer: {}", balance_after); + } self .emit( @@ -158,10 +230,15 @@ mod streaming { recipient: ContractAddress, start_time: u64, end_time: u64, - ) -> (u128, u128) { + ) -> (u128, u128, bool, ContractAddress) { let key: (ContractAddress, u64, u64) = (recipient, start_time, end_time); - let (currently_claimable, total_amount): (u128, u128) = self.streams.read(key); - (currently_claimable, total_amount) + let ( + currently_claimable, total_amount, is_minting, token_address + ): (u128, u128, bool, ContractAddress) = + self + .streams + .read(key); + (currently_claimable, total_amount, is_minting, token_address) } } } diff --git a/src/treasury.cairo b/src/treasury.cairo index e628faed..623daf4e 100644 --- a/src/treasury.cairo +++ b/src/treasury.cairo @@ -3,6 +3,7 @@ use starknet::ContractAddress; #[starknet::interface] trait ITreasury { + // Existing Treasury functions fn send_tokens_to_address( ref self: TContractState, receiver: ContractAddress, @@ -35,6 +36,26 @@ trait ITreasury { fn withdraw_from_nostra_lending_pool( ref self: TContractState, nostraToken: ContractAddress, amount: u256 ); + + // Streaming functions + fn add_new_stream( + ref self: TContractState, + recipient: ContractAddress, + start_time: u64, + end_time: u64, + total_amount: u128, + is_minting: bool, + token_address: ContractAddress + ); + fn claim_stream( + ref self: TContractState, recipient: ContractAddress, start_time: u64, end_time: u64, + ); + fn cancel_stream( + ref self: TContractState, recipient: ContractAddress, start_time: u64, end_time: u64, + ); + fn get_stream_info( + self: @TContractState, recipient: ContractAddress, start_time: u64, end_time: u64, + ) -> (u128, u128, bool, ContractAddress); } #[starknet::contract] @@ -43,7 +64,11 @@ mod Treasury { use core::starknet::event::EventEmitter; use core::traits::TryInto; use konoha::airdrop::{IAirdropDispatcher, IAirdropDispatcherTrait}; + + use konoha::contract::Governance; + use konoha::contract::{IGovernanceDispatcher, IGovernanceDispatcherTrait}; use konoha::traits::{IERC20Dispatcher, IERC20DispatcherTrait}; + use konoha::traits::{IGovernanceTokenDispatcher, IGovernanceTokenDispatcherTrait}; use konoha::treasury_types::carmine::{IAMMDispatcher, IAMMDispatcherTrait}; use konoha::treasury_types::nostra::interface::{ INostraInterestToken, INostraInterestTokenDispatcher, INostraInterestTokenDispatcherTrait @@ -55,7 +80,10 @@ mod Treasury { use openzeppelin::access::ownable::interface::IOwnableTwoStep; use openzeppelin::upgrades::interface::IUpgradeable; use openzeppelin::upgrades::upgradeable::UpgradeableComponent; - use starknet::{ContractAddress, get_caller_address, get_contract_address, ClassHash}; + + use starknet::{ + ContractAddress, get_caller_address, get_contract_address, ClassHash, get_block_timestamp + }; use super::{OptionType}; component!(path: OwnableComponent, storage: ownable, event: OwnableEvent); component!(path: UpgradeableComponent, storage: upgradeable, event: UpgradeableEvent); @@ -73,8 +101,35 @@ mod Treasury { #[substorage(v0)] ownable: OwnableComponent::Storage, #[substorage(v0)] - upgradeable: UpgradeableComponent::Storage + upgradeable: UpgradeableComponent::Storage, + streams: LegacyMap<(ContractAddress, u64, u64), (u128, u128, bool, ContractAddress)>, + } + + #[derive(starknet::Event, Drop)] + struct StreamCreated { + recipient: ContractAddress, + start_time: u64, + end_time: u64, + total_amount: u128, + token_address: ContractAddress, + } + + #[derive(starknet::Event, Drop)] + struct StreamClaimed { + recipient: ContractAddress, + start_time: u64, + end_time: u64, + amount_claimed: u128, + } + + #[derive(starknet::Event, Drop)] + struct StreamCanceled { + recipient: ContractAddress, + start_time: u64, + end_time: u64, + reclaimed_amount: u128, } + #[derive(starknet::Event, Drop)] struct TokenSent { receiver: ContractAddress, @@ -128,10 +183,12 @@ mod Treasury { amount: u256 } - #[event] #[derive(Drop, starknet::Event)] enum Event { + StreamCreated: StreamCreated, + StreamClaimed: StreamClaimed, + StreamCanceled: StreamCanceled, TokenSent: TokenSent, AMMAddressUpdated: AMMAddressUpdated, LiquidityProvided: LiquidityProvided, @@ -350,8 +407,133 @@ mod Treasury { LiquidityWithdrawnFromNostraLendingPool { nostra_token: nostraToken, amount } ); } + + //starts a new stream through the treasury. + //mint = true to mint CRM, false to transfer treasury tokens to recipient + fn add_new_stream( + ref self: ContractState, + recipient: ContractAddress, + start_time: u64, + end_time: u64, + total_amount: u128, + is_minting: bool, + token_address: ContractAddress, + ) { + self.ownable.assert_only_owner(); + let key = (recipient, start_time, end_time); + assert(start_time < end_time, 'starts first'); + + let claimable_amount = 0; + self.streams.write(key, (claimable_amount, total_amount, is_minting, token_address)); + + self + .emit( + StreamCreated { recipient, start_time, end_time, total_amount, token_address } + ); + } + //claim tokens from a time between the start and end of the stream + fn claim_stream( + ref self: ContractState, recipient: ContractAddress, start_time: u64, end_time: u64, + ) { + let current_time = get_block_timestamp(); + let key = (recipient, start_time, end_time); + let (already_claimed, total_amount, is_minting, token_address) = self.streams.read(key); + + assert(current_time > start_time, 'stream has not started'); + + let elapsed_time = if current_time > end_time { + end_time - start_time + } else { + current_time - start_time + }; + let stream_duration = end_time - start_time; + + let currently_claimable = (total_amount * elapsed_time.into() / stream_duration.into()); + let amount_to_claim = currently_claimable - already_claimed; + + assert(amount_to_claim > 0, 'nothing to claim'); + + self.streams.write(key, (currently_claimable, total_amount, is_minting, token_address)); + + if is_minting { + let governance_dispatcher = IGovernanceDispatcher { + contract_address: get_contract_address() + }; + let governance_token_address = governance_dispatcher.get_governance_token_address(); + let governance_token_dispatcher = IGovernanceTokenDispatcher { + contract_address: governance_token_address, + }; + + // Mint the tokens directly to the recipient + governance_token_dispatcher.mint(recipient, amount_to_claim.into()); + } else { + //transfer treasury tokens + self.internal_send_tokens(recipient, amount_to_claim.into(), token_address); + } + + self + .emit( + StreamClaimed { + recipient, start_time, end_time, amount_claimed: amount_to_claim + } + ); + } + //end the stream and mint the rest of the CRM tokens to user + fn cancel_stream( + ref self: ContractState, recipient: ContractAddress, start_time: u64, end_time: u64, + ) { + self.ownable.assert_only_owner(); + let key = (recipient, start_time, end_time); + let (already_claimed, total_amount, is_minting, token_address) = self.streams.read(key); + + let to_distribute = total_amount - already_claimed; + + self.streams.write(key, (0, 0, is_minting, token_address)); + + if is_minting { + let governance_dispatcher = IGovernanceDispatcher { + contract_address: get_contract_address() + }; + let governance_token_address = governance_dispatcher.get_governance_token_address(); + let governance_token_dispatcher = IGovernanceTokenDispatcher { + contract_address: governance_token_address, + }; + //mints the rest of the tokens to the recipient + governance_token_dispatcher.mint(recipient, to_distribute.into()); + } else { + self.internal_send_tokens(recipient, to_distribute.into(), token_address); + } + + self + .emit( + StreamCanceled { + recipient, start_time, end_time, reclaimed_amount: to_distribute + } + ); + } + + fn get_stream_info( + self: @ContractState, recipient: ContractAddress, start_time: u64, end_time: u64, + ) -> (u128, u128, bool, ContractAddress) { + let key = (recipient, start_time, end_time); + self.streams.read(key) + } } + #[generate_trait] + impl InternalFunctions of InternalTrait { + fn internal_send_tokens( + ref self: ContractState, + receiver: ContractAddress, + amount: u256, + token_addr: ContractAddress + ) { + let token: IERC20Dispatcher = IERC20Dispatcher { contract_address: token_addr }; + assert(token.balance_of(get_contract_address()) >= amount, Errors::INSUFFICIENT_FUNDS); + token.transfer(receiver, amount); + self.emit(TokenSent { receiver, token_addr, amount }); + } + } #[abi(embed_v0)] impl UpgradeableImpl of IUpgradeable { diff --git a/tests/lib.cairo b/tests/lib.cairo index a7df965e..67a531af 100644 --- a/tests/lib.cairo +++ b/tests/lib.cairo @@ -4,7 +4,6 @@ mod proposals_tests; mod setup; mod staking_tests; mod test_storage_pack; -mod test_streaming; mod test_treasury; mod upgrades_tests; mod vesting; diff --git a/tests/test_streaming.cairo b/tests/test_streaming.cairo deleted file mode 100644 index 7b5b5df8..00000000 --- a/tests/test_streaming.cairo +++ /dev/null @@ -1,173 +0,0 @@ -use array::ArrayTrait; -use core::option::OptionTrait; -use core::result::ResultTrait; -use core::traits::TryInto; -use debug::PrintTrait; - -use konoha::contract::Governance; -use konoha::contract::{IGovernanceDispatcher, IGovernanceDispatcherTrait}; -use konoha::streaming::{IStreamingDispatcher, IStreamingDispatcherTrait, IStreaming}; -use konoha::traits::{IGovernanceTokenDispatcher, IGovernanceTokenDispatcherTrait}; -use konoha::vesting::{IVestingDispatcher, IVestingDispatcherTrait, IVesting}; - -use openzeppelin::token::erc20::interface::{IERC20Dispatcher, IERC20DispatcherTrait}; -use snforge_std::{ - BlockId, declare, ContractClassTrait, ContractClass, start_prank, start_warp, CheatTarget, - prank, CheatSpan -}; -use starknet::{get_block_timestamp, get_caller_address, get_contract_address, ContractAddress}; - -use super::setup::{deploy_governance_and_both_tokens}; - -fn start_stream(gov: ContractAddress) { - prank(CheatTarget::One(gov), gov, CheatSpan::TargetCalls(4)); - let streaming = IStreamingDispatcher { contract_address: gov }; - streaming.add_new_stream(0x2.try_into().unwrap(), 100, 200, 100000); -} - -//passing! -#[test] -fn test_add_new_stream() { - let (gov, _, _) = deploy_governance_and_both_tokens(); - start_stream(gov.contract_address); - let streaming = IStreamingDispatcher { contract_address: gov.contract_address }; - - let recipient = 0x2.try_into().unwrap(); - let start_time: u64 = 100; - let end_time: u64 = 200; - let total_amount: u128 = 100000; - - streaming.add_new_stream(recipient, start_time, end_time, total_amount); - //let key = (get_caller_address(), recipient, end_time, start_time); - - let (claimed_amount, stored_total_amount) = streaming - .get_stream_info(recipient, start_time, end_time,); - - assert_eq!(recipient, 0x2.try_into().unwrap(), "Incorrect streamer addr"); - assert_eq!(start_time, 100, "Incorrect start time"); - assert_eq!(end_time, 200, "Incorrect end time"); - assert_eq!(claimed_amount, 0, "Incorrect claimed amount after stream creation"); - assert_eq!(stored_total_amount, 100000, "Incorrect total amount stored"); -} - -//passing! -#[test] -#[should_panic(expected: ('starts first',))] -fn test_valid_stream_time() { - let (gov, _, _) = deploy_governance_and_both_tokens(); - start_stream(gov.contract_address); - let streaming = IStreamingDispatcher { contract_address: gov.contract_address }; - - let recipient = 0x2.try_into().unwrap(); - let start_time: u64 = 200; - let end_time: u64 = 100; - let total_amount: u128 = 100000; - - streaming.add_new_stream(recipient, start_time, end_time, total_amount); -} - -//passing! -#[test] -#[should_panic(expected: ('nothing to claim',))] -fn test_claimed_amount() { - let (gov, _, _) = deploy_governance_and_both_tokens(); - start_stream(gov.contract_address); - let streaming = IStreamingDispatcher { contract_address: gov.contract_address }; - - let recipient = 0x2.try_into().unwrap(); - let start_time: u64 = 100; - let end_time: u64 = 200; - let total_amount: u128 = 0; - streaming.add_new_stream(recipient, start_time, end_time, total_amount); - - start_warp(CheatTarget::One(gov.contract_address), 150); - //shouldn't have anything to claim - streaming.claim_stream(recipient, start_time, end_time); -} - -//passing! -#[test] -#[should_panic(expected: ('stream has not started',))] -fn test_stream_started() { - let (gov, _, _) = deploy_governance_and_both_tokens(); - start_stream(gov.contract_address); - let streaming = IStreamingDispatcher { contract_address: gov.contract_address }; - - let recipient = 0x2.try_into().unwrap(); - let start_time: u64 = 100; - let end_time: u64 = 200; - let total_amount: u128 = 100000; - streaming.add_new_stream(recipient, start_time, end_time, total_amount); - start_warp(CheatTarget::One(gov.contract_address), 50); // before of stream - - streaming.claim_stream(recipient, start_time, end_time); -} - -#[test] -fn test_claim_stream() { - let (gov, _, _) = deploy_governance_and_both_tokens(); - start_stream(gov.contract_address); - let streaming = IStreamingDispatcher { contract_address: gov.contract_address }; - - let recipient = 0x2.try_into().unwrap(); - let start_time: u64 = 100; - let end_time: u64 = 200; - let total_amount: u128 = 100000; - - streaming.add_new_stream(recipient, start_time, end_time, total_amount); - let (claimable_amount, total_amount) = streaming - .get_stream_info(recipient, start_time, end_time,); - start_warp(CheatTarget::One(gov.contract_address), 150); - - streaming.claim_stream(recipient, start_time, end_time); - - let expected_claimed_amount = (100000 * 50 / 100); //should be 50% since middle of stream - assert_eq!(total_amount, 100000, "Incorrect total amount after claiming the stream"); - assert_eq!(claimable_amount, 0, "Incorrect claimed amount after claiming the stream"); - - let self_dsp = IGovernanceDispatcher { contract_address: gov.contract_address }; - let token_address = self_dsp.get_governance_token_address(); - let erc20 = IERC20Dispatcher { contract_address: token_address }; - - let balance = erc20.balance_of(recipient); - - assert_eq!( - balance, expected_claimed_amount, "Balance should match the expected claimed amount" - ); -} - -#[test] -fn test_cancel_stream() { - let (gov, _, _) = deploy_governance_and_both_tokens(); - start_stream(gov.contract_address); - let streaming = IStreamingDispatcher { contract_address: gov.contract_address }; - - let recipient = 0x2.try_into().unwrap(); - let start_time: u64 = 100; - let end_time: u64 = 200; - let total_amount: u128 = 100000; - - streaming.add_new_stream(recipient, start_time, end_time, total_amount); - - start_warp(CheatTarget::One(gov.contract_address), 150); - - //test cancel_stream - streaming.cancel_stream(recipient, start_time, end_time); - - let (claimed_amount, stored_total_amount) = streaming - .get_stream_info(recipient, start_time, end_time); - - assert_eq!(claimed_amount, 0, "Claimed amount should be 0 after canceling the stream"); - assert_eq!(stored_total_amount, 0, "Total amount should be 0 after canceling the stream"); - - let unclaimed_amount: u256 = total_amount.into() - claimed_amount.into(); //100000 - let self_dsp = IGovernanceDispatcher { contract_address: gov.contract_address }; - let token_address = self_dsp.get_governance_token_address(); - let erc20 = IERC20Dispatcher { contract_address: token_address }; - - // Check the balance of the streamer (caller address) with ERC_20, I couldnt use balance_of - let balance = erc20.balance_of(get_caller_address()); - - assert_eq!(unclaimed_amount.into(), 100000, "Unclaimed amount should be reclaimed correctly"); - assert_eq!(balance, 0, "balance"); -} diff --git a/tests/test_treasury.cairo b/tests/test_treasury.cairo index caa57b33..859f6687 100644 --- a/tests/test_treasury.cairo +++ b/tests/test_treasury.cairo @@ -6,7 +6,10 @@ use core::result::ResultTrait; use core::serde::Serde; use core::traits::{TryInto, Into}; use debug::PrintTrait; +use konoha::contract::Governance; +use konoha::contract::{IGovernanceDispatcher, IGovernanceDispatcherTrait}; use konoha::traits::{IERC20Dispatcher, IERC20DispatcherTrait}; +use konoha::traits::{IGovernanceTokenDispatcher, IGovernanceTokenDispatcherTrait}; use konoha::treasury::{ITreasuryDispatcher, ITreasuryDispatcherTrait}; use konoha::treasury_types::carmine::{IAMMDispatcher, IAMMDispatcherTrait}; use konoha::treasury_types::zklend::interfaces::{IMarketDispatcher, IMarketDispatcherTrait}; @@ -15,9 +18,12 @@ use openzeppelin::access::ownable::interface::{ }; use openzeppelin::upgrades::interface::{IUpgradeableDispatcher, IUpgradeableDispatcherTrait}; use snforge_std::{ - BlockId, declare, ContractClassTrait, ContractClass, prank, CheatSpan, CheatTarget, roll + BlockId, declare, ContractClassTrait, ContractClass, prank, CheatSpan, CheatTarget, roll, + start_warp }; use starknet::{ContractAddress, get_block_number, ClassHash}; +use super::setup::{deploy_governance_and_both_tokens}; + mod testStorage { use core::traits::TryInto; use starknet::ContractAddress; @@ -63,6 +69,158 @@ fn get_important_addresses() -> ( ); } +#[test] +#[fork("MAINNET")] +fn test_transfer_stream() { + let (gov_contract_address, _, treasury_contract_address, _) = get_important_addresses(); + let eth_addr: ContractAddress = + 0x049d36570d4e46f48e99674bd3fcc84644ddd6b96f7c741b1562b82f9e004dc7 + .try_into() + .unwrap(); + let random_whale: ContractAddress = + 0x4267ae838da77a52384283f3321a0746557023d24cb823115d2da5c8c4f1a42 + .try_into() + .unwrap(); + let eth_dispatcher = IERC20Dispatcher { contract_address: eth_addr }; + let treasury_dispatcher = ITreasuryDispatcher { contract_address: treasury_contract_address }; + + // Transfer eth from random whale to Treasury contract + prank(CheatTarget::One(eth_addr), random_whale, CheatSpan::TargetCalls(1)); + let deposit_amt = 200000; + eth_dispatcher.transfer(treasury_contract_address, deposit_amt); + let initial_treasury_bal = eth_dispatcher.balance_of(treasury_contract_address); + assert(initial_treasury_bal >= deposit_amt, 'eth bal too low'); + println!("Initial Treasury Balance {}:", initial_treasury_bal); + + let recipient: ContractAddress = 0x2.try_into().unwrap(); + let start_time: u64 = 100; + let end_time: u64 = 200; + let total_amount: u128 = 2000; + let is_minting: bool = false; + let token_address: ContractAddress = + 0x049d36570d4e46f48e99674bd3fcc84644ddd6b96f7c741b1562b82f9e004dc7 + .try_into() + .unwrap(); + + // Add new stream + prank( + CheatTarget::One(treasury_contract_address), gov_contract_address, CheatSpan::TargetCalls(1) + ); + treasury_dispatcher + .add_new_stream(recipient, start_time, end_time, total_amount, is_minting, token_address); + + let (claimed, total, is_minting_stored, token_stored) = treasury_dispatcher + .get_stream_info(recipient, start_time, end_time); + assert(claimed == 0, 'Initial claim should be 0'); + assert(total == total_amount, 'Total amount mismatch'); + assert(is_minting_stored == is_minting, 'Is minting mismatch'); + assert(token_stored == token_address, 'Token address mismatch'); + + start_warp(CheatTarget::One(treasury_contract_address), 150); + prank(CheatTarget::One(treasury_contract_address), recipient, CheatSpan::TargetCalls(1)); + treasury_dispatcher.claim_stream(recipient, start_time, end_time); + + // Verify claimed amount + let (claimed_midway, _, _, _) = treasury_dispatcher + .get_stream_info(recipient, start_time, end_time); + let recipient_balance_after = eth_dispatcher.balance_of(recipient); + + assert(claimed_midway > 0 && claimed_midway < total_amount, 'Midway claim incorrect'); + println!("recipient balance after {}:", recipient_balance_after); //is 1000 + + let recipient_balance_before_cancel = eth_dispatcher.balance_of(recipient); + prank( + CheatTarget::One(treasury_contract_address), gov_contract_address, CheatSpan::TargetCalls(1) + ); + treasury_dispatcher.cancel_stream(recipient, start_time, end_time); + + // Verify stream is cancelled and remaining tokens transferred + let (claimed_final, total_final, _, _) = treasury_dispatcher + .get_stream_info(recipient, start_time, end_time); + let recipient_balance_after_cancel = eth_dispatcher.balance_of(recipient); + assert(claimed_final == 0 && total_final == 0, 'Stream not properly cancelled'); + assert(recipient_balance_after_cancel > recipient_balance_before_cancel, 'Remaining tokens'); +} + +#[test] +fn test_mint_stream() { + let (_, _, _, _) = get_important_addresses(); // Assuming treasury address is not used + let (gov, _, _) = deploy_governance_and_both_tokens(); + let recipient = 0x2.try_into().unwrap(); + let start_time: u64 = 100; + let end_time: u64 = 200; + let total_amount: u128 = 100000; + let expected_claimed_amount = (total_amount * 50 / 100); // 50% claimed since middle of stream + let is_minting = true; + + let treasury_dispatcher = ITreasuryDispatcher { contract_address: gov.contract_address }; + let erc20 = IERC20Dispatcher { + contract_address: IGovernanceDispatcher { contract_address: gov.contract_address } + .get_governance_token_address() + }; + let governance_token_address = IGovernanceDispatcher { contract_address: gov.contract_address } + .get_governance_token_address(); + + // Test with minting + treasury_dispatcher + .add_new_stream( + recipient, start_time, end_time, total_amount, is_minting, governance_token_address + ); + + let (claimed_amount, stored_total_amount, is_minting, token_address) = treasury_dispatcher + .get_stream_info(recipient, start_time, end_time); + assert_eq!(claimed_amount, 0, "Incorrect claimed amount after stream creation (minting)"); + assert_eq!(stored_total_amount, total_amount, "Incorrect total amount stored (minting)"); + assert_eq!(is_minting, true, "Stream should be set to minting"); + println!("Streaming: {:?} token", token_address); + + // Test claim with minting + start_warp(CheatTarget::One(gov.contract_address), 150); + treasury_dispatcher.claim_stream(recipient, start_time, end_time); + + let (claimed_amount, stored_total_amount, _, _) = treasury_dispatcher + .get_stream_info(recipient, start_time, end_time); + assert_eq!( + stored_total_amount, + total_amount, + "Incorrect total amount after claiming the stream (minting)" + ); + assert_eq!( + claimed_amount, + expected_claimed_amount, + "Incorrect claimed amount after claiming the stream (minting)" + ); + + let recipient_balance = erc20.balance_of(recipient); + assert_eq!( + recipient_balance, + expected_claimed_amount.into(), + "Recipient balance should match the expected claimed amount (minting)" + ); //50000 + + println!("recipinet balance before cancel: {}", recipient_balance); + // Test cancel with minting + treasury_dispatcher.cancel_stream(recipient, start_time, end_time); + + let (cancelled_claimed_amount, stored_total_amount, _, _) = treasury_dispatcher + .get_stream_info(recipient, start_time, end_time); + assert_eq!( + cancelled_claimed_amount, + 0, + "Claimed amount should be 0 after canceling the stream (minting)" + ); + assert_eq!( + stored_total_amount, 0, "Total amount should be 0 after canceling the stream (minting)" + ); + + // Verify recipient's balance after cancellation + let recipient_balance_after_cancel = erc20.balance_of(recipient); + assert_eq!( + recipient_balance_after_cancel, + total_amount.into(), + "Recipient balance should match the claimed amount after cancellation (minting)" + ); +} #[test] #[fork("SEPOLIA")]