Skip to main content

Write unit tests

It's time to write unit tests for your first Equito App. In this section, you will learn how to write comprehensive unit tests for the PingPong contract featured in the previous section of this tutorial.

Setting up the environment​

Foundry provides a built-in testing framework that allows you to write and run tests for your contracts. You can create test files in the test/ folder and use the Test contract from the forge-std library to define test cases.

Create a new file PingPongTest.t.sol in the test/ folder with the following content:

pragma solidity ^0.8.23;

import {Test, console} from "forge-std/Test.sol";
import {Vm} from "forge-std/Vm.sol";
import {PingPong} from "@equito/src/examples/PingPong.sol";
import {Router} from "@equito/src/Router.sol";
import {MockVerifier} from "@equito/test/mock/MockVerifier.sol";
import {MockEquitoFees} from "@equito/test/mock/MockEquitoFees.sol";
import {bytes64, EquitoMessage, EquitoMessageLibrary} from "@equito/src/libraries/EquitoMessageLibrary.sol";
import {Errors} from "@equito/src/libraries/Errors.sol";

contract PingPongTest is Test {
    // Test cases go here...
}

Initialize the contracts​

Before writing test cases, you need to initialize the contracts that you want to test and their dependencies. In this case, you need to initialize the PingPong contract and its dependencies such as the Router, MockVerifier, and MockEquitoFees contracts.
The Router will be deployed with chain selector 0, and the PingPong contract will have a couple of simulated peers with chain selectors 1 and 2, meaning that these addresses do not match any deployed contract.

We start from defining which contracts we want to deploy, some useful addresses, and the events that we expect to be emitted.

contract PingPongTest is Test {
    Router router;
    MockVerifier verifier;
    MockEquitoFees fees;
    PingPong pingPong;

    address sender = address(0xa0);
    address equitoAddress = address(0xe0);
    address peer1 = address(0x01);
    address peer2 = address(0x02);

    event PingSent(
        uint256 indexed destinationChainSelector,
        bytes32 messageHash
    );
    event PingReceived(
        uint256 indexed sourceChainSelector,
        bytes32 messageHash
    );
    event PongReceived(
        uint256 indexed sourceChainSelector,
        bytes32 messageHash
    );
    event MessageSendRequested(
      EquitoMessage message,
      bytes messageData
    );

    // Test cases go here...
}

Then we define the setUp function, which Foundry will call before running each test case.

contract PingPongTest is Test {
    // Previous code...

    function setUp() public override {
        vm.prank(sender);

        // Deploy the Equito contracts
        verifier = new MockVerifier();
        fees = new MockEquitoFees();
        router = new Router(
            0, // local chainSelector
            address(verifier),
            address(fees),
            EquitoMessageLibrary.addressToBytes64(equitoAddress)
        );

        // Deploy the PingPong Contract
        pingPong = new PingPong(address(router));

        // Set the peers
        uint256[] memory chainSelectors = new uint256[](2);
        chainSelectors[0] = 1;
        chainSelectors[1] = 2;

        bytes64[] memory addresses = new bytes64[](2);
        addresses[0] = EquitoMessageLibrary.addressToBytes64(peer1);
        addresses[1] = EquitoMessageLibrary.addressToBytes64(peer2);

        pingPong.setPeers(chainSelectors, addresses);

        // Fund the sender
        vm.deal(address(sender), 10 ether);
    }

    // Test cases go here...
}

Write test cases​

Test sending a ping​

The first test case is to send a ping message from the sender to a peer. We expect the PingSent event to be emitted.

Since the event emits the destinationChainSelector and the messageHash, we need to calculate the messageHash based on the expected message. We should always remember to pay the right amount of fees to the router contract.

function testSendPing() public {
    // Useful variables for the message construction
    uint256 destinationChainSelector = 2;
    string memory pingMessage = "Ping!";
    bytes memory messageData = abi.encode("ping", pingMessage);

    // Calculate the fee
    uint256 fee = router.getFee(address(pingPong));

    // We expect a PingSent event, containing the destination chain and message hash
    vm.prank(sender);
    vm.expectEmit(true, true, true, true);
    emit PingSent(
        destinationChainSelector,
        keccak256(
            abi.encode(
                EquitoMessage({
                    blockNumber: 1,
                    sourceChainSelector: 0,
                    sender: EquitoMessageLibrary.addressToBytes64(
                        address(pingPong)
                    ),
                    destinationChainSelector: destinationChainSelector,
                    receiver: EquitoMessageLibrary.addressToBytes64(peer2),
                    hashedData: keccak256(messageData)
                })
            )
        )
    );

    // Send the ping message, paying the fee
    pingPong.sendPing{value: fee}(destinationChainSelector, pingMessage);
}

Test receiving a ping and sending a pong​

The second test case is to receive a ping message from a peer and send a pong message back. We expect the emission of a PingReceived and a MessageSendRequest, the event that Equito uses to send messages.

function testReceivePingAndSendPong() public {
    // Forge a valid ping message
    string memory pingMessage = "Equito";
    bytes memory messageData = abi.encode("ping", pingMessage);
    EquitoMessage memory message = EquitoMessage({
        blockNumber: 1,
        sourceChainSelector: 2,
        sender: EquitoMessageLibrary.addressToBytes64(peer2),
        destinationChainSelector: 0,
        receiver: EquitoMessageLibrary.addressToBytes64(address(pingPong)),
        hashedData: keccak256(messageData)
    });

    // Calculate the fee and the current balance of the fee contract
    uint256 feeContractBalance = address(fees).balance;
    uint256 fee = router.getFee(address(pingPong));

    // We expect a PingReceived event and a MessageSendRequested event
    vm.expectEmit(address(pingPong));
    emit PingReceived(2, keccak256(abi.encode(message)));

    vm.expectEmit(address(router));
    emit MessageSendRequested(
        EquitoMessage({
            blockNumber: 1,
            sourceChainSelector: 0,
            sender: EquitoMessageLibrary.addressToBytes64(
                address(pingPong)
            ),
            destinationChainSelector: 2,
            receiver: EquitoMessageLibrary.addressToBytes64(peer2),
            hashedData: keccak256(abi.encode("pong", pingMessage))
        }),
        abi.encode("pong", pingMessage)
    );

    // Receive the message and send a pong back
    // Make sure to pay the fee for the pong message
    vm.prank(address(sender));
    router.deliverAndExecuteMessage{value: fee}(
        message,
        messageData,
        0,
        abi.encode(1) // The MockVerifier returns true for non-empty proofs
    );

    // Ensure the fee was paid correctly
    assertEq(address(fees).balance, feeContractBalance + fee);
}

Test receiving a pong​

To finalize the correct path, we test receiving a pong message from a peer. We expect the emission of a PongReceived event.

function testReceivePong() public {
    string memory pongMessage = "Pong!";

    // Forge a valid pong message
    bytes memory messageData = abi.encode("pong", pongMessage);
    EquitoMessage memory message = EquitoMessage({
        blockNumber: 1,
        sourceChainSelector: 1,
        sender: EquitoMessageLibrary.addressToBytes64(peer1),
        destinationChainSelector: 2,
        receiver: EquitoMessageLibrary.addressToBytes64(peer2),
        hashedData: keccak256(messageData)
    });

    // Pretend to be the router
    vm.prank(address(router));

    // We expect a PongReceived event
    vm.expectEmit(true, true, true, true);
    emit PongReceived(1, keccak256(abi.encode(message)));

    // Receive the pong message
    pingPong.receiveMessage(message, messageData);
}

Test receiving an invalid message​

We start testing the wrong paths by testing the reception of an invalid message, where we expect the emission of an InvalidMessageType error.

function testInvalidMessageType() public {
    string memory invalidMessage = "Invalid";

    // Forge a message with invalid type
    // Supported types are "ping" and "pong"
    bytes memory messageData = abi.encode("invalid", invalidMessage);
    EquitoMessage memory message = EquitoMessage({
        blockNumber: 1,
        sourceChainSelector: 1,
        sender: EquitoMessageLibrary.addressToBytes64(peer1),
        destinationChainSelector: 2,
        receiver: EquitoMessageLibrary.addressToBytes64(peer2),
        hashedData: keccak256(messageData)
    });

    // Pretend to be the router
    vm.prank(address(router));

    // We expect an InvalidMessageType error
    vm.expectRevert(PingPong.InvalidMessageType.selector);

    // Receive the invalid message
    pingPong.receiveMessage(message, messageData);
}

Test receiving a message from an unauthorized sender​

We can finally test that the behavior of the contract is correct when receiving a message from an invalid sender, i.e. when it is not a peer. We expect the emission of an InvalidMessageSender error.

function testInvalidSender() public {
    string memory pongMessage = "message";

    // Forge a valid message from an invalid sender
    bytes memory messageData = abi.encode("pong", pongMessage);
    EquitoMessage memory message = EquitoMessage({
        blockNumber: 1,
        sourceChainSelector: 1,
        sender: EquitoMessageLibrary.addressToBytes64(address(0x00)),
        destinationChainSelector: 0,
        receiver: EquitoMessageLibrary.addressToBytes64(address(pingPong)),
        hashedData: keccak256(messageData)
    });

    // Pretend to be the router
    vm.prank(address(router));

    vm.expectRevert(Errors.InvalidMessageSender.selector);
    pingPong.receiveMessage(message, messageData);
}

Running the tests​

To run the test with Foundry, you can simply execute the following command:

forge test