Skip to content

Helper functions#

Wake testing framework provides a set of helper functions to make testing easier.

ABI encoding and decoding#

The Abi class provides functions to encode and decode data according to the ABI specification.

Abi.encode#

Abi.encode encodes a list of values given a list of types. It returns bytes:

from wake.testing import Abi, Address

Abi.encode(['uint8', 'address'], [0xff, Address(0)])

Abi.encode_packed#

Abi.encode_packed encodes a list of values given a list of types. It returns bytes:

from wake.testing import Abi

Abi.encode_packed(['bytes', 'string'], [b'abc', 'def'])

Abi.encode_with_selector#

Abi.encode_with_selector encodes a list of values and a selector given a list of types and the selector. It returns bytes:

from wake.testing import Abi
from pytypes.contracts.Counter import Counter

Abi.encode_with_selector(Counter.setCount.selector, ['uint256'], [0xff])

Abi.encode_with_signature#

Abi.encode_with_signature encodes a list of values and a selector given a list of types and a signature. It returns bytes:

from wake.testing import Abi

Abi.encode_with_signature("setCount(uint256)", ['uint256'], [0xff])

Warning

The signature string must conform to the ABI specification. The common mistakes are:

  • uint or int used instead of uint256 or int256,
  • return type specified,
  • spaces in the signature string.

Abi.encode_call#

Abi.encode_call encodes a list of values and a selector given a reference to a function. It returns bytes:

from wake.testing import Abi
from pytypes.contracts.Counter import Counter

Abi.encode_call(Counter.setCount, [0xff])

Abi.decode#

Abi.decode decodes a bytes object given a list of types. It returns a list of values:

from wake.testing import Abi

Abi.decode(['uint8', 'address'],
           b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\xff\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00')

Keccak-256#

The keccak256 function computes the Keccak-256 hash of a bytes object:

from wake.testing import keccak256

keccak256(b'abc')

Computing CREATE and CREATE2 address#

In some cases, it may be useful to compute the address of a contract before it is deployed. Wake testing framework provides three functions to do so.

get_create_address#

get_create_address computes the address of a contract deployed in a transaction or in a contract using the CREATE opcode. It accepts a deployer (Account, Address or a hex string address) and its nonce.

from wake.testing import Account, get_create_address

deployer = Account(1)
get_create_address(deployer, deployer.nonce)

get_create2_address_from_code#

get_create2_address_from_code computes the address of a contract deployed using the CREATE2 opcode. It accepts a deployer (Account, Address or a hex string address), a salt and the contract creation code.

from wake.testing import Account, get_create2_address_from_code
from wake.testing.fuzzing import random_bytes
from pytypes.contracts.Counter import Counter

get_create2_address_from_code(
    Account(1),
    random_bytes(32),
    Counter.get_creation_code()
)

get_create2_address_from_hash#

get_create2_address_from_hash computes the address of a contract deployed using the CREATE2 opcode. It accepts a deployer (Account, Address or a hex string address), a salt and the hash of the contract creation code.

from wake.testing import Account, get_create2_address_from_hash, keccak256
from wake.testing.fuzzing import random_bytes
from pytypes.contracts.Counter import Counter

get_create2_address_from_hash(
    Account(1),
    random_bytes(32),
    keccak256(Counter.get_creation_code())
)

Get logic contract from proxy#

get_logic_contract returns the logic contract Account from a proxy Account. If the input account is not a proxy, it returns the input account.

from wake.testing import Account, get_logic_contract

usdc_proxy = Account("0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48")
usdc_logic = get_logic_contract(usdc_proxy)

Read & write storage variable#

read_storage_variable and write_storage_variable read and write storage variables of a contract. They accept a contract Account and a variable name. The keys argument is used to specify the array indexes, mapping keys and struct member names needed to access the variable.

Info

Reading and writing arrays is done using lists, while reading and writing structs is done using dictionaries.

Reading and writing whole mappings is not supported.

If the provided contract is a proxy, the variable definition is searched in the logic contract and the proxy storage is used. This behavior can be overridden by setting the storage_layout_contract argument. In this case, the variable definition is searched in the provided storage_layout_contract.

from wake.testing import Account, Address, read_storage_variable, write_storage_variable

usdc_proxy = Account("0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48")
write_storage_variable(usdc_proxy, "balances", 1000, keys=[Address(1)])
assert read_storage_variable(usdc_proxy, "balances", keys=[Address(1)]) == 1000

ERC-20 mint and burn#

mint_erc20 and burn_erc20 mint and burn ERC-20 tokens. They detect the totalSupply and balances variables using heuristics and may not work for all contracts. Optionally, balance_slot and total_supply_slot arguments can be used to specify the storage slot where the balance of the given account and the total supply are stored.

from wake.testing import Account, mint_erc20, burn_erc20
from pytypes.contracts.IERC20 import IERC20

usdc_proxy = IERC20("0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48")
burn_erc20(usdc_proxy, Account(1), usdc_proxy.balanceOf(Account(1)))
mint_erc20(usdc_proxy, Account(1), 1000)
assert usdc_proxy.balanceOf(Account(1)) == 1000

Decorators#

on_revert#

on_revert is a decorator that simplifies handling of revert exceptions. It accepts a callback function that will be called if the decorated function reverts.

from wake.testing import *


def revert_handler(e: TransactionRevertedError):
    if e.tx is not None:
        print(e.tx.call_trace)
        print(e.tx.console_logs)


@chain.connect()
@on_revert(revert_handler)
def test_reverts():
    ...