Foundry с 0. Часть 29

Продолжаем совершенствоваться с инвариант тестами и сегодня поговорим о таком подвиде тестов, как Handler Based Tests.

Помните, в предыдущем посте мы делали тесты контракта WETH, в котором в случайном порядке вызывались вообще все функции, что приводило к большому количеству откатов. Мы сделали вывод, что такой тест нам не подходит и нужно провести его более прицельно. Как раз в этом нам и поможет дополнительный контракт помощник - Handler.

Я приведу его полный пример вместе с обновленным тестом и мы разберем его "по кирпичикам". Итак:

import {CommonBase} from "forge-std/Base.sol";
import {StdCheats} from "forge-std/StdCheats.sol";
import {StdUtils} from "forge-std/StdUtils.sol";

contract Handler is CommonBase, StdCheats, StdUtils {
    WETH private weth;
    uint256 public wethBalance;
    uint256 public numCalls;

    constructor(WETH _weth) {
        weth = _weth;
    }

    receive() external payable {}

    function sendToFallback(uint256 amount) public {
        amount = bound(amount, 0, address(this).balance);
        wethBalance += amount;
        numCalls += 1;

        (bool ok,) = address(weth).call{value: amount}("");
        require(ok, "sendToFallback failed");
    }

    function deposit(uint256 amount) public {
        amount = bound(amount, 0, address(this).balance);
        wethBalance += amount;
        numCalls += 1;

        weth.deposit{value: amount}();
    }

    function withdraw(uint256 amount) public {
        amount = bound(amount, 0, weth.balanceOf(address(this)));
        wethBalance -= amount;
        numCalls += 1;

        weth.withdraw(amount);
    }

    function fail() external {
        revert("fail");
    }
}

И контракт нашего теста:

contract WETH_Handler_Based_Invariant_Tests is Test {
    WETH public weth;
    Handler public handler;

    function setUp() public {
        weth = new WETH();
        handler = new Handler(weth);

        deal(address(handler), 100 * 1e18);
        targetContract(address(handler));

        bytes4[] memory selectors = new bytes4[](3);
        selectors[0] = Handler.deposit.selector;
        selectors[1] = Handler.withdraw.selector;
        selectors[2] = Handler.sendToFallback.selector;

        targetSelector(
            FuzzSelector({addr: address(handler), selectors: selectors})
        );
    }

    function invariant_eth_balance() public {
        assertGe(address(weth).balance, handler.wethBalance());
        console.log("handler num calls", handler.numCalls());
    }
}

Разберем наш Handler.

Handler - это некая смесь обычного контракта с тестовым, благодаря новым подключаемым библиотекам: CommonBase, StdCheats и StdUtils. Именно они позволяют использовать читкоды в Handler.

Итак, тут мы создаем переменные для нашего контракта WETH, который определим в конструкторе, а также wethBalance и numCalls, которые помогут написать хорошие тесты.

В данном примере мы будем тестировать три функции контракта WETH: deposit, withdraw и fallback, для этих целей мы создадим в Handler свои функции, которые будут вызывать функции из WETH.

function sendToFallback(uint256 amount) public {}
function deposit(uint256 amount) public {}
function withdraw(uint256 amount) public {}

Как вы можете заметить мы применим небольшие элементы фаззинга и позволим Foundry также отправлять различные суммы. При этом мы хотим, чтобы эти суммы были в некоторых рамках баланса нашего Handler. Ведь, если будут вообще все возможные суммы, то количество ревертов сильно возрастет. Ограничить можно с помощью уже пройдённого нами читкода:

amount = bound(amount, 0, weth.balanceOf(address(this)));

Также не забываем записывать в переменные отправляемые значение и увеличивать количество вызовов на 1.

Теперь контракт теста.

Для начала мы создаем переменные для объектов наших контрактов и прописываем функцию setUp(), которая играет определяющую роль.

weth = new WETH();
handler = new Handler(weth);

Инициализируем наши контракты.

deal(address(handler), 100 * 1e18);

Пополняем баланс Handler на 100 Эфиров, чтобы можно было что-то отправлять на WETH.