aztec-packages
aztec-packages copied to clipboard
TXE: `with_impersonated` / `with_set_contract_address`
Consider methods like:
pub fn check_public_balance(token_contract_address: AztecAddress, user: AztecAddress, expected_amount: Field) {
let current_contract_address = get_contract_address();
cheatcodes::set_contract_address(token_contract_address);
let block_number = get_block_number();
let slot = Token::storage().bets.slot;
let address_slot = derive_storage_slot_in_map(slot, user);
let amount: U128 = storage_read(token_contract_address, address_slot, block_number);
assert(amount.to_field() == expected_amount);
cheatcodes::set_contract_address(current_contract_address);
}
you have to remember to reset set_contract_address
(same for impersonate_account
)
Should we perhaps have something like this:
impl TestEnvironment {
fn with_impersonated<T, Env>(self, account: AztecAddress, f: fn[Env]() -> T) -> T {
let old_impersonated = self.contract_address();
self.impersonate(account);
let result = f();
self.impersonate(old_impersonated);
result
}
}