diff --git a/snforge_std/src/cheatcodes.cairo b/snforge_std/src/cheatcodes.cairo index 1b3665bc9d..ccd82b1f80 100644 --- a/snforge_std/src/cheatcodes.cairo +++ b/snforge_std/src/cheatcodes.cairo @@ -6,6 +6,7 @@ pub mod l1_handler; pub mod contract_class; pub mod storage; pub mod execution_info; +pub mod erc20; pub mod message_to_l1; pub mod generate_random_felt; @@ -43,7 +44,7 @@ pub fn test_address() -> ContractAddress { /// - `ret_data` - data to return by the function `function_selector` /// - `n_times` - number of calls to mock the function for pub fn mock_call<T, impl TSerde: core::serde::Serde<T>, impl TDestruct: Destruct<T>>( - contract_address: ContractAddress, function_selector: felt252, ret_data: T, n_times: u32 + contract_address: ContractAddress, function_selector: felt252, ret_data: T, n_times: u32, ) { assert!(n_times > 0, "cannot mock_call 0 times, n_times argument must be greater than 0"); @@ -68,7 +69,7 @@ pub fn mock_call<T, impl TSerde: core::serde::Serde<T>, impl TDestruct: Destruct /// macro) /// - `ret_data` - data to be returned by the function pub fn start_mock_call<T, impl TSerde: core::serde::Serde<T>, impl TDestruct: Destruct<T>>( - contract_address: ContractAddress, function_selector: felt252, ret_data: T + contract_address: ContractAddress, function_selector: felt252, ret_data: T, ) { let contract_address_felt: felt252 = contract_address.into(); let mut inputs = array![contract_address_felt, function_selector]; @@ -91,7 +92,7 @@ pub fn start_mock_call<T, impl TSerde: core::serde::Serde<T>, impl TDestruct: De pub fn stop_mock_call(contract_address: ContractAddress, function_selector: felt252) { let contract_address_felt: felt252 = contract_address.into(); execute_cheatcode_and_deserialize::< - 'stop_mock_call', () + 'stop_mock_call', (), >(array![contract_address_felt, function_selector].span()); } @@ -111,9 +112,9 @@ pub enum ReplaceBytecodeError { /// Returns `Result::Ok` if the replacement succeeded, and a `ReplaceBytecodeError` with appropriate /// error type otherwise pub fn replace_bytecode( - contract: ContractAddress, new_class: ClassHash + contract: ContractAddress, new_class: ClassHash, ) -> Result<(), ReplaceBytecodeError> { execute_cheatcode_and_deserialize::< - 'replace_bytecode' + 'replace_bytecode', >(array![contract.into(), new_class.into()].span()) } diff --git a/snforge_std/src/cheatcodes/erc20.cairo b/snforge_std/src/cheatcodes/erc20.cairo new file mode 100644 index 0000000000..714b37d4bb --- /dev/null +++ b/snforge_std/src/cheatcodes/erc20.cairo @@ -0,0 +1,49 @@ +use starknet::{ContractAddress}; +use snforge_std::cheatcodes::storage::{map_entry_address, store}; + +const STRK_CONTRACT_ADDRESS: felt252 = + 0x04718f5a0fc34cc1af16a1cdee98ffb20c31f5cd61d6ab07201858f4287c938d; + +#[derive(Drop, Serde, Copy)] +pub struct CustomToken { + pub contract_address: ContractAddress, + pub balances_variable_selector: felt252, +} + +#[derive(Drop, Copy, Clone)] +pub enum Token { + STRK, + Custom: CustomToken, +} + +#[generate_trait] +pub impl TokenImpl of TokenTrait { + fn contract_address(self: Token) -> ContractAddress { + match self { + Token::STRK => STRK_CONTRACT_ADDRESS.try_into().unwrap(), + Token::Custom(CustomToken { contract_address, .. }) => contract_address, + } + } + + fn balances_variable_selector(self: Token) -> felt252 { + match self { + Token::STRK => selector!("ERC20_balances"), + Token::Custom(CustomToken { balances_variable_selector, + .., }) => balances_variable_selector, + } + } +} + +/// Sets the balance of `token`` for `target`` contract to `new_balance` +/// - `target` - address of the contract, which balance you want to modify +/// - `new_balance` - new balance value +/// - `token` - token for which the balance is being set +pub fn set_balance(target: ContractAddress, new_balance: u256, token: Token) { + let balance_low_address = map_entry_address( + token.balances_variable_selector(), [target.into()].span(), + ); + let balance_high_address = balance_low_address + 1; + + store(token.contract_address(), balance_high_address, array![new_balance.low.into()].span()); + store(token.contract_address(), balance_low_address, array![new_balance.high.into()].span()); +} \ No newline at end of file diff --git a/snforge_std/src/lib.cairo b/snforge_std/src/lib.cairo index e62e9e6006..f742e21763 100644 --- a/snforge_std/src/lib.cairo +++ b/snforge_std/src/lib.cairo @@ -26,6 +26,12 @@ pub use cheatcodes::storage::store; pub use cheatcodes::storage::load; pub use cheatcodes::storage::map_entry_address; +pub use cheatcodes::erc20::set_balance; +pub use cheatcodes::erc20::Token; +pub use cheatcodes::erc20::TokenImpl; +pub use cheatcodes::erc20::TokenTrait; +pub use cheatcodes::erc20::CustomToken; + pub use cheatcodes::CheatSpan; pub use cheatcodes::ReplaceBytecodeError; pub use cheatcodes::test_address;