Zero to Hero in Foundry - Part 4: Fuzz & Invariant Tests

Recap

Awesome! I see you’re back!

Part 3 was a banger with some good old testing fundamentals. We wrote a pretty simple bank contract through which users could deposit and withdraw funds, as well as check their balances. Deployed the same using anvil and forge and wrote fundamental tests for making sure that all of the functionalities worked as intended

Please read it if you haven't already before continuing

Foundry: Zero to Hero - Part 3

Wanna learn Web3 through live and interactive challenges? You’ll love and find that Web3Compass is the best!!

Today's Outcome

Topic of focus: Fuzz and Invariant Tests

Code base - TokenSwap ↗️

• Write 2 ERC20 token contracts - TokenA and TokenB

• Write a TokenSwap.sol contract to swap our ERC20s

• Let users add liquidity for both tokens in our swap contract

• Let users swap our tokens making sure that there’s enough balance in user’s wallet and our contract pool

• And the most important bit - Write extensive Fuzz and Invariant tests

SHALL WE?!

---

The What ⁉️

What is Fuzz Testing? 🥴

Fuzz testing is what happens when we give our code an espresso, a blindfold, and a keyboard. It’s an automated technique where we let a "fuzzer" throw absolute nonsense at our functions to see what works (or, more likely, what breaks).

We all have that one test we forget to write. What if the input is 115792089237316195423570985008687907853269984665640564039457584007913129639935? What if it's a negative number in a uint? Instead of losing sleep over it, we let Foundry's fuzzer do our dirty work. We just add parameters to a test function, and Foundry unleashes a relentless barrage of random inputs, hoping to make our code cry. When it fails, the fuzzer proudly presents the input that caused the chaos.

Analogy: Imagine our function is a nightclub bouncer. Normal testing is showing the bouncer a few valid IDs. Fuzz testing is sending a thousand aliens, ghosts, and cartoon characters with crayon-drawn IDs to the front door, just to see if one of them somehow gets in.

What is Invariant Testing? 🏛️

If fuzz testing is poking a single function with a stick, invariant testing is shaking the entire building to see if the foundations are solid.

An invariant is a sacred, unbreakable rule of our contract. It’s a "pinky promise" that must always be true, no matter what happens. For example:

• The total supply of our precious token doesn't just invent itself out of thin air.

• The contract’s piggy bank should never be empty if it still owes people money.

Invariant testing lets Foundry's fuzzer go completely wild on our entire contract. It will call any function it wants, in any order, like a baby who found the TV remote. Its only goal is to break one of our sacred rules. If it manages to shatter a promise, it tells us exactly how it did it, so we can fix our flawed logic.

Analogy: Let's say our contract is a smoothie stall. Our main invariant is: "The number of bananas we have must always equal to the number of bananas we started with, minus the number we've sold." An invariant test is like unleashing a horde of hyperactive squirrels to randomly buy smoothies, sell us more bananas, and try to knock over the stall for a full day. At the end, we check if our banana count is still correct. If not, we've got a squirrelly bug.

---

The How 🤔

Let’s go ahead and see how we can implement these tests. For this we need to have a contract. Time to implement our Tokens and TokenSwap contracts

Writing the Contracts 📝

TokenA.sol

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.30;

import "@openzeppelin/contracts/token/ERC20/ERC20.sol";

contract TokenA is ERC20 {
    constructor() ERC20("TokenA", "TKA") {
        _mint(msg.sender, 10000e18);
    }
}

This creates a basic ERC20 token called TokenA using OpenZeppelin’s secure ERC20 implementation.

• Name and Symbol: The token is named "TokenA" and its symbol is "TKA".

• Initial Supply: When the contract is deployed, it mints 10,000 tokens (with 18 decimals) to the deployer’s address.

• ERC20 Standard: By inheriting from OpenZeppelin’s ERC20, TokenA supports all standard ERC20 features like transferring tokens, checking balances, and approving allowances.

Similar to this let’s write TokenB.sol for our TokenB supply

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.30;

import "@openzeppelin/contracts/token/ERC20/ERC20.sol";

contract TokenB is ERC20 {
    constructor() ERC20("TokenB", "TKB") {
        _mint(msg.sender, 10000e18);
    }
}

OKAYYY!!! There are our token contracts written and ready to be tested. Time to write our swapping contract

TokenSwap.sol

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.30;

import "@openzeppelin/contracts/token/ERC20/IERC20.sol"; // Importing the ERC20 interface

contract TokenSwap {
    IERC20 public tokenA; // TokenA contract instance
    IERC20 public tokenB; // TokenB contract instance

    // Constructor to set the token addresses
    constructor(address _tokenA, address _tokenB) {
        tokenA = IERC20(_tokenA);
        tokenB = IERC20(_tokenB);
    }

    // Function to swap TokenA for TokenB
    function swapAforB(uint256 amount) external {
        require(amount > 0, "Amounts must be greater than zero");
        require(tokenA.balanceOf(msg.sender) >= amount, "Insufficient TokenA balance in Wallet");
        require(tokenB.balanceOf(address(this)) >= amount, "Insufficient TokenB balance in Pool");
        tokenA.transferFrom(msg.sender, address(this), amount);
        tokenB.transfer(msg.sender, amount);
    }

    // Function to swap TokenB for TokenA
    function swapBforA(uint256 amount) external {
        require(amount > 0, "Amounts must be greater than zero");
        require(tokenB.balanceOf(msg.sender) >= amount, "Insufficient TokenB balance in Wallet");
        require(tokenA.balanceOf(address(this)) >= amount, "Insufficient TokenA balance in Pool");
        tokenB.transferFrom(msg.sender, address(this), amount);
        tokenA.transfer(msg.sender, amount);
    }

    // Function to add liquidity for TokenA
    function addTokenALiquidity(uint256 amount) external {
        require(amount > 0, "Amount must be greater than zero");
        require(tokenA.balanceOf(msg.sender) >= amount, "Insufficient TokenA balance");
        tokenA.transferFrom(msg.sender, address(this), amount);
    }

    // Function to add liquidity for TokenB
    function addTokenBLiquidity(uint256 amount) external {
        require(amount > 0, "Amount must be greater than zero");
        require(tokenB.balanceOf(msg.sender) >= amount, "Insufficient TokenB balance");
        tokenB.transferFrom(msg.sender, address(this), amount);
    }

}

The contract lets users swap between two ERC20 tokens (TokenA and TokenB) and add liquidity to the pool.

• Token Setup:
  The contract is initialized with the addresses of TokenA and TokenB. It uses the standard ERC20 interface to interact with both tokens.

• Swapping Tokens:

  • swapAforB(amount): Users can swap TokenA for TokenB. The contract checks that the user has enough TokenA and the pool has enough TokenB. It then transfers TokenA from the user to the contract and TokenB from the contract to the user.

  • swapBforA(amount): Users can swap TokenB for TokenA, following the same logic in reverse.

• Adding Liquidity:

  • addTokenALiquidity(amount): Users can deposit TokenA into the pool, increasing the contract’s TokenA balance.

  • addTokenBLiquidity(amount): Users can deposit TokenB into the pool, increasing the contract’s TokenB balance.

• Security Checks:
  Each function checks that the amount is greater than zero and that the sender has enough balance before proceeding.

Fuzz & Invariant Tests 🧪

Our contracts are in place

• TokenA.sol

• TokenB.sol

• TokenSwap.sol

They’re completed and ready to be tested inside and out. No point lollygagging anymore, let’s jump straight into out Fuzz and Invariant test code

TokenSwap.t.sol

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.30;

import "forge-std/Test.sol";

import "../src/TokenSwap.sol";
import "../src/TokenA.sol";
import "../src/TokenB.sol";

contract TokenSwapTest is Test {
    TokenA tokenA;
    TokenB tokenB;
    TokenSwap tokenSwap;

    address user = address(0xABCD); // User address

    /// @dev This function is called before each test
    /// @notice This function sets up the token contracts and initial balances
    function setUp() public {
        tokenA = new TokenA();
        tokenB = new TokenB();
        tokenSwap = new TokenSwap(address(tokenA), address(tokenB));

        // Mint tokens to user and contract for liquidity
        tokenA.transfer(user, 1000e18);
        tokenB.transfer(user, 1000e18);
        tokenB.transfer(address(tokenSwap), 1000e18);
        tokenA.transfer(address(tokenSwap), 1000e18);

        // Approve the TokenSwap contract to spend user tokens
        vm.startPrank(user);
        tokenA.approve(address(tokenSwap), 1000e18);
        tokenB.approve(address(tokenSwap), 1000e18);
        vm.stopPrank();
    }

    /// @dev Fuzz test for addTokenALiquidity
    /// @notice This function tests the addTokenALiquidity function with various input amounts
    function testFuzz_addTokenALiquidity(uint256 amount) public {
        amount = bound(amount, 1, 1000e18); // Bound the amount to a valid range

        // Start test for addTokenALiquidity
        vm.startPrank(user);
        tokenA.approve(address(tokenSwap), amount);
        tokenSwap.addTokenALiquidity(amount);
        // Check liquidity added
        assertEq(tokenA.balanceOf(address(tokenSwap)), 1000e18 + amount);
        assertEq(tokenA.balanceOf(user), 1000e18 - amount);
        vm.stopPrank();
    }

    /// @dev Fuzz test for addTokenBLiquidity
    /// @notice This function tests the addTokenBLiquidity function with various input amounts
    function testFuzz_addTokenBLiquidity(uint256 amount) public {
        amount = bound(amount, 1, 1000e18);

        // Start test for addTokenBLiquidity
        vm.startPrank(user);
        tokenB.approve(address(tokenSwap), amount);
        tokenSwap.addTokenBLiquidity(amount);
        // Check liquidity added
        assertEq(tokenB.balanceOf(address(tokenSwap)), 1000e18 + amount);
        assertEq(tokenB.balanceOf(user), 1000e18 - amount);
        vm.stopPrank();
    }

    /// @dev Fuzz test for swapAforB
    /// @notice This function tests the swapAforB function with various input amounts
    function testFuzz_swapAforB(uint256 amount) public {
        amount = bound(amount, 1, 1000e18);
        vm.assume(tokenA.balanceOf(user) >= amount); // User must have enough TokenA
        vm.assume(tokenB.balanceOf(address(tokenSwap)) >= amount); // TokenSwap must have enough TokenB

        // Start test for swapAforB
        vm.startPrank(user);
        tokenA.approve(address(tokenSwap), amount);
        uint256 userTokenBBefore = tokenB.balanceOf(user);
        uint256 contractTokenBBefore = tokenB.balanceOf(address(tokenSwap));
        tokenSwap.swapAforB(amount);
        // Check balances after swap
        assertEq(tokenA.balanceOf(user), 1000e18 - amount);
        assertEq(tokenB.balanceOf(user), userTokenBBefore + amount);
        assertEq(tokenB.balanceOf(address(tokenSwap)), contractTokenBBefore - amount);
        vm.stopPrank();
    }

    /// @dev Fuzz test for swapBforA
    /// @notice This function tests the swapBforA function with various input amounts
    function testFuzz_swapBforA(uint256 amount) public {
        amount = bound(amount, 1, 1000e18);
        vm.assume(tokenB.balanceOf(user) >= amount);
        vm.assume(tokenA.balanceOf(address(tokenSwap)) >= amount);

        // Start test for swapBforA
        vm.startPrank(user);
        tokenB.approve(address(tokenSwap), amount);
        uint256 userTokenABefore = tokenA.balanceOf(user);
        uint256 contractTokenABefore = tokenA.balanceOf(address(tokenSwap));
        tokenSwap.swapBforA(amount);
        // Check balances after swap
        assertEq(tokenB.balanceOf(user), 1000e18 - amount);
        assertEq(tokenA.balanceOf(user), userTokenABefore + amount);
        assertEq(tokenA.balanceOf(address(tokenSwap)), contractTokenABefore - amount);
        vm.stopPrank();
    }

    // Invariant: total supply of TokenA and TokenB never changes
    /// @notice This function checks that the total supply of both tokens remains constant
    function invariant_totalSupplyConstant() public view {
        assertEq(tokenA.totalSupply(), 10000e18);
        assertEq(tokenB.totalSupply(), 10000e18);
    }

    // Invariant: sum of TokenA balances between user and contract is always <= 10000e18
    /// @notice This function checks that the sum of TokenA balances between user and contract is always <= 10000e18
    function invariant_tokenABalanceSum() public view {
        uint256 sum = tokenA.balanceOf(user) + tokenA.balanceOf(address(tokenSwap));
        assertLe(sum, 10000e18);
    }

    // Invariant: sum of TokenB balances between user and contract is always <= 10000e18
    /// @notice This function checks that the sum of TokenB balances between user and contract is always <= 10000e18
    function invariant_tokenBBalanceSum() public view {
        uint256 sum = tokenB.balanceOf(user) + tokenB.balanceOf(address(tokenSwap));
        assertLe(sum, 10000e18);
    }
}

Here the test contract rigorously checks the behavior of the TokenSwap smart contract. Here’s what it does:

• Setup:
  Before each test, it deploys fresh instances of TokenA, TokenB, and TokenSwap. It gives both the user and the TokenSwap contract 1,000 tokens each for liquidity and sets up the necessary approvals.

• Fuzz Tests:
  These tests use random amounts to check the swap and liquidity functions:

  • testFuzz_addTokenALiquidity and testFuzz_addTokenBLiquidity verify that users can add liquidity in various amounts and that balances update correctly.

  • testFuzz_swapAforB and testFuzz_swapBforA check that swapping tokens works for any valid amount, ensuring balances change as expected.

• Invariant Tests:
  These tests ensure that certain properties always hold, no matter what sequence of actions is performed:

  • invariant_totalSupplyConstant checks that the total supply of both tokens never changes.

  • invariant_tokenABalanceSum and invariant_tokenBBalanceSum ensure that the sum of user and contract balances for each token never exceeds the original supply.

• Testing Tools:
  The contract uses Foundry’s vm.startPrank to simulate actions from the user’s address and bound to keep fuzzed amounts within a valid range.

ANDDD, let’s test! Use forge to start testing our swap contract

forge coverage

There we have our test results. BOO YAA!! Amazing work!

We can do this to any type of contract and make sure our contracts are bullet proof. Maybe try writing a token Borrowing contract and hit it with Fuzz and Invariant tests. It’s always great to learn by doing. We now have the power of testing contracts rigorously, and like someone wise once told

---

And that's a Wrap! (For now 😬)

By now we should have

• Our token and swap contracts written

• Can add liquidity to our swap contract and swap our tokens

• We have written proper & extensive Fuzz and Invariant tests for it

• Checked our test coverage and got close to 100%

What's Next?

Part 5, we'll be

• Writing a batch token sender contract

• Focusing majorly on getting Gas reports and optimizations

Don't miss it! I'll see you soon! 😎