Add set_balance Cairo implementation

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())
 }

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());
+}

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;