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.