Writing Invariant Tests with Narya SDK

Introduction

Invariant tests are a type of test used to check if a particular property holds true in all possible program states. In the case of smart contracts, these tests are essential to ensure that the contract behaves correctly and does not have any security vulnerabilities.

This document provides a step-by-step guide on how to write your first invariant test using the Narya SDK.

Example Production Contract

Consider the following production contract that we want to test:

pragma solidity ^0.8.0;

contract Vault {
    mapping(address => uint256) public balances;

    // Allows users to put some tokens on contract address
    function deposit() public payable {
        balances[msg.sender] = balances[msg.sender] + msg.value;
    }

    // Allows users to get tokens out
    function withdraw() public {
        if (balances[msg.sender] > 0) {
            uint256 balance = balances[msg.sender];
            (bool result, ) = msg.sender.call{value: balance}("");
            require(result, "Failed to withdraw");
            balances[msg.sender] = 0;
        }
    }
}

Invariant Statement

Invariants are statements that we want to ensure hold true for all possible program states. In this case, we want to ensure that the user can withdraw the same amount of tokens that they deposited before. Thus, we can phrase the invariant statement as follows:

If user A deposits N tokens to the contract, then no matter what other users do, the contract balance should always be greater than or equal to N.

If we can find a way to violate this invariant, our Vault contract has a serious vulnerability.

Invariant Test

Now let's see how we can translate this statement into an invariant test using the Narya SDK. In this example, we will use the Narya SDK with Solidity version 0.8.0.

pragma solidity ^0.8.0;

import {Vault} from "contracts/Vault.sol";
import {PTest} from "@narya-ai/contracts/PTest.sol";

// A test contract, needs to extends `PTest`
// Contract name could be anything
contract FundLossTest is PTest {
  Vault vault;

  address user = makeAddr("user");
  address agent;

  function setUp() external {
    // Step 1: Deploy target contract
    vault = new Vault();

    // Step 2: Grant user account 1 eth, and deposit it into the vault
    hoax(user, 1 ether); // Next call will be sent from `user` account
    vault.deposit{value: 1 ether}();

    // Step 3: Grant agent account 1 eth, and deposit it into the vault
    agent = getAgent(); // Initialize the field
    hoax(agent, 1 ether); // Next call will be sent from `agent` account
    vault.deposit{value: 1 ether}();
  }

  // Invariant assertion function.
  // Function name needs to be started with `invariant`.
  // Prefer to have no side effect (view function), but this is not required.
  function invariantFundLoss() public view {
    // Invariant: Contract balance should always be >= 1 ether,
    // since the user takes no action after setUp() is called.
    assert(address(vault).balance >= 1 ether);
  }
}

Let's break it down:

First, we import Vault and PTest contracts. Vault is the contract we are testing, and PTest is extended by the test smart contract for our engine to work with it.

Next, we declare a setUp function that contains three steps:

Deploy the Vault smart contract:
```
vault = new Vault();
```

Deposit 1 ether on behalf of the user:

hoax(user, 1 ether); // Next call will be sent from `user` account
vault.deposit{value: 1 ether}();

Deposit 1 ether on behalf of the agent:

agent = getAgent(); // Initialize the field
hoax(agent, 1 ether); // Next call will be sent from `agent` account
vault.deposit{value: 1 ether}();

In the above code, the agent variable represents a hacker bot controlled by the engine, whose sole purpose is to try to break the invariants defined. You may want to give the agent some setup, depending on your test cases. You can always get the address by calling getAgent(). It’s a standard address. In our example, we gave it 1 ether, and then deposited the 1 ether to the vault on its behalf.

Finally, we define the invariant:

Vault contract balance should always be >= 1 ether

assert(address(vault).balance >= 1 ether);

You can find this example project at NaryaAi/example-vault.

All test types: Test Types
More test examples: NaryaAI/narya-examples
Cheat codes (test utilities):
- Foundry cheat codes
- Narya cheat codes

Next Step

Next, you could use the Narya CLI tool to quickly validate the test locally. It doesn't have the same power as our cloud engine, but it's a quick way to make sure the test has the right format and is compilable. Alternatively, you could go ahead and run the tests on the cloud. See Run Tests on the Cloud for more information.

Narya Docs