aztec-packages icon indicating copy to clipboard operation
aztec-packages copied to clipboard

TXE: `with_impersonated` / `with_set_contract_address`

Open rahul-kothari opened this issue 6 months ago • 0 comments

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

rahul-kothari avatar Aug 12 '24 16:08 rahul-kothari