Skip to content

Commit

Permalink
test: inherit "Base_Test" in "BaseHandler"
Browse files Browse the repository at this point in the history 
test: rename "LockupHandlerStore" to "LockupHandlerStorage"
test: use "excludeSender" cheatcode for core contracts
test: use "excludeSender" cheatcode for handlers contracts
in the invariants tests
test: remove unnecessary function from "LockupHandlerStorage"
test: rename "flashLoan" contract to "flashLoanContract"
test: rename "flashLoanFunction" function to "flashLoan"
test: rename "store" to "_storage"
chore: remove unused imports
  • Loading branch information
@andreivladbrg
andreivladbrg committed Feb 6, 2023
1 parent 9db6ddb commit fa5c007
Show file tree
Hide file tree
Showing 14 changed files with 122 additions and 143 deletions.
2 changes: 1 addition & 1 deletion test/Base.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -66,8 +66,8 @@ abstract contract Base_Test is Assertions, Constants, Calculations, Utils, StdCh
ISablierV2Comptroller internal comptroller;
IERC20 internal dai = new ERC20("Dai Stablecoin", "DAI");
ISablierV2LockupLinear internal linear;
NonCompliantERC20 internal nonCompliantAsset = new NonCompliantERC20("Non-Compliant ERC-20 Asset", "NCT", 18);
ISablierV2LockupPro internal pro;
NonCompliantERC20 internal nonCompliantAsset = new NonCompliantERC20("Non-Compliant ERC-20 Asset", "NCT", 18);

/*//////////////////////////////////////////////////////////////////////////
CONSTRUCTOR
Expand Down
7 changes: 6 additions & 1 deletion test/invariant/Invariant.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@ import { InvariantTest as ForgeInvariantTest } from "forge-std/InvariantTest.sol

import { SablierV2Comptroller } from "src/SablierV2Comptroller.sol";

import { GoodFlashLoanReceiver } from "../shared/mockups/flash-loan/GoodFlashLoanReceiver.t.sol";
import { Base_Test } from "../Base.t.sol";
import { ComptrollerHandler } from "./handlers/ComptrollerHandler.t.sol";

Expand Down Expand Up @@ -34,6 +33,12 @@ abstract contract Invariant_Test is Base_Test, ForgeInvariantTest {
// Target only the comptroller handler for invariant testing (to avoid getting reverts).
targetContract(address(comptrollerHandler));

// Exclude the comptroller, linear and pro for being the `msg.sender`.
excludeSender(address(comptroller));

// Exclude the comptroller handler for being the `msg.sender`.
excludeSender(address(comptrollerHandler));

// Label the base contracts.
vm.label({ account: address(comptroller), newLabel: "Comptroller" });
vm.label({ account: address(comptrollerHandler), newLabel: "ComptrollerHandler" });
Expand Down
18 changes: 3 additions & 15 deletions test/invariant/handlers/BaseHandler.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,23 +8,11 @@ import { Calculations } from "../../shared/helpers/Calculations.t.sol";
import { Constants } from "../../shared/helpers/Constants.t.sol";
import { Utils } from "../../shared/helpers/Utils.t.sol";

import { Base_Test } from "../../Base.t.sol";

/// @title BaseHandler
/// @notice Base contract with common logic needed by all handler contracts.
abstract contract BaseHandler is Constants, Calculations, Utils, StdCheats {
/*//////////////////////////////////////////////////////////////////////////
CONSTANTS
//////////////////////////////////////////////////////////////////////////*/

/// @dev The address of the HEVM contract.
address internal constant HEVM_ADDRESS = address(uint160(uint256(keccak256("hevm cheat code"))));

/*//////////////////////////////////////////////////////////////////////////
TEST CONTRACTS
//////////////////////////////////////////////////////////////////////////*/

/// @dev An instance of the HEVM.
Vm internal constant vm = Vm(HEVM_ADDRESS);

abstract contract BaseHandler is Base_Test {
/*//////////////////////////////////////////////////////////////////////////
TEST VARIABLES
//////////////////////////////////////////////////////////////////////////*/
Expand Down
6 changes: 0 additions & 6 deletions test/invariant/handlers/ComptrollerHandler.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,12 +12,6 @@ import { BaseHandler } from "./BaseHandler.t.sol";
/// @dev This contract and not {SablierV2Comptroller} is exposed to Foundry for invariant testing. The point is
/// to bound and restrict the inputs that get passed to the real-world contract to avoid getting reverts.
contract ComptrollerHandler is BaseHandler {
/*//////////////////////////////////////////////////////////////////////////
TEST CONTRACTS
//////////////////////////////////////////////////////////////////////////*/

ISablierV2Comptroller public comptroller;

/*//////////////////////////////////////////////////////////////////////////
CONSTRUCTOR
//////////////////////////////////////////////////////////////////////////*/
Expand Down
23 changes: 9 additions & 14 deletions test/invariant/handlers/FlashLoanHandler.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,20 +8,18 @@ import { UD60x18 } from "@prb/math/UD60x18.sol";
import { SablierV2FlashLoan } from "src/abstracts/SablierV2FlashLoan.sol";
import { ISablierV2Comptroller } from "src/interfaces/ISablierV2Comptroller.sol";

import { GoodFlashLoanReceiver } from "../../shared/mockups/flash-loan/GoodFlashLoanReceiver.t.sol";
import { BaseHandler } from "./BaseHandler.t.sol";

/// @title FlashLoanHandler
/// @dev This contract and not {SablierV2FlashLoan} is exposed to Foundry for invariant testing. The point is
/// to bound and restrict the inputs that get passed to the real-world contract to avoid getting reverts.
contract FlashLoanHandler is BaseHandler {
/*//////////////////////////////////////////////////////////////////////////
PUBLIC TEST CONTRACTS
TEST CONTRACTS
//////////////////////////////////////////////////////////////////////////*/

IERC20 public asset;
ISablierV2Comptroller public comptroller;
SablierV2FlashLoan public flashLoan;
SablierV2FlashLoan public flashLoanContract;
IERC3156FlashBorrower internal receiver;

/*//////////////////////////////////////////////////////////////////////////
Expand All @@ -31,47 +29,44 @@ contract FlashLoanHandler is BaseHandler {
constructor(
IERC20 asset_,
ISablierV2Comptroller comptroller_,
SablierV2FlashLoan flashLoan_,
SablierV2FlashLoan flashLoanContract_,
IERC3156FlashBorrower receiver_
) {
asset = asset_;
comptroller = comptroller_;
flashLoan = flashLoan_;
receiver = receiver_;
flashLoanContract = flashLoanContract_;
receiver_ = receiver;

This comment has been minimized.

Sign in to view

Copy link
@PaulRBerg

PaulRBerg Feb 10, 2023

Member

@andreivladbrg I just noticed that you introduced a bug here.

I will fix this in #322.
}

/*//////////////////////////////////////////////////////////////////////////
SABLIER-V2-FLASH-LOAN
//////////////////////////////////////////////////////////////////////////*/

function flashLoanFunction(uint128 amount) external instrument("flashLoan") {
function flashLoan(uint128 amount) external instrument("flashLoan") {
// Only supported ERC-20 assets can be flash loaned.
bool isFlashLoanable = comptroller.isFlashLoanable(asset);
if (!isFlashLoanable) {
return;
}

// The flash fee must be less than or equal to type(uint128).max
uint256 fee = flashLoan.flashFee(address(asset), amount);
uint256 fee = flashLoanContract.flashFee(address(asset), amount);
if (fee > type(uint128).max) {
return;
}

// Mint the flash loan amount to the contract.
deal({ token: address(asset), to: address(flashLoan), give: amount });
deal({ token: address(asset), to: address(flashLoanContract), give: amount });

// Mint the flash fee to the receiver so that they can repay the flash loan.
deal({ token: address(asset), to: address(receiver), give: fee });

// Execute the flash loan.
bool response = flashLoan.flashLoan({
flashLoanContract.flashLoan({
receiver: receiver,
asset: address(asset),
amount: amount,
data: bytes("Some Data")
});

// Silence the compiler warning.
response;
}
}
28 changes: 14 additions & 14 deletions test/invariant/handlers/LockupHandler.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,18 +7,18 @@ import { ISablierV2Lockup } from "src/interfaces/ISablierV2Lockup.sol";
import { Broker, Lockup } from "src/types/DataTypes.sol";

import { BaseHandler } from "./BaseHandler.t.sol";
import { LockupHandlerStore } from "./LockupHandlerStore.t.sol";
import { LockupHandlerStorage } from "./LockupHandlerStorage.t.sol";

/// @title LockupHandler
/// @dev Common handler logic between {SablierV2LockupLinear} and {SablierV2LockupPro}.
abstract contract LockupHandler is BaseHandler {
/*//////////////////////////////////////////////////////////////////////////
PUBLIC TEST CONTRACTS
TEST CONTRACTS
//////////////////////////////////////////////////////////////////////////*/

IERC20 public asset;
ISablierV2Lockup public lockup;
LockupHandlerStore public store;
LockupHandlerStorage public _storage;

/*//////////////////////////////////////////////////////////////////////////
PRIVATE TEST VARIABLES
Expand All @@ -32,10 +32,10 @@ abstract contract LockupHandler is BaseHandler {
CONSTRUCTOR
//////////////////////////////////////////////////////////////////////////*/

constructor(IERC20 asset_, ISablierV2Lockup lockup_, LockupHandlerStore store_) {
constructor(IERC20 asset_, ISablierV2Lockup lockup_, LockupHandlerStorage _storage_) {
asset = asset_;
lockup = lockup_;
store = store_;
_storage = _storage_;
}

/*//////////////////////////////////////////////////////////////////////////
Expand All @@ -50,24 +50,24 @@ abstract contract LockupHandler is BaseHandler {
}

modifier useFuzzedStreamRecipient(uint256 streamIndexSeed) {
uint256 lastStreamId = store.lastStreamId();
uint256 lastStreamId = _storage.lastStreamId();
if (lastStreamId == 0) {
return;
}
currentStreamId = store.streamIds(bound(streamIndexSeed, 0, lastStreamId - 1));
currentRecipient = store.streamIdsToRecipients(currentStreamId);
currentStreamId = _storage.streamIds(bound(streamIndexSeed, 0, lastStreamId - 1));
currentRecipient = _storage.streamIdsToRecipients(currentStreamId);
vm.startPrank(currentRecipient);
_;
vm.stopPrank();
}

modifier useFuzzedStreamSender(uint256 streamIndexSeed) {
uint256 lastStreamId = store.lastStreamId();
uint256 lastStreamId = _storage.lastStreamId();
if (lastStreamId == 0) {
return;
}
currentStreamId = store.streamIds(bound(streamIndexSeed, 0, lastStreamId - 1));
currentSender = store.streamIdsToSenders(currentStreamId);
currentStreamId = _storage.streamIds(bound(streamIndexSeed, 0, lastStreamId - 1));
currentSender = _storage.streamIdsToSenders(currentStreamId);
vm.startPrank(currentSender);
_;
vm.stopPrank();
Expand All @@ -93,7 +93,7 @@ abstract contract LockupHandler is BaseHandler {
lockup.burn(currentStreamId);

// Set the recipient associated with this stream to the zero address.
store.updateRecipient(currentStreamId, address(0));
_storage.updateRecipient(currentStreamId, address(0));
}

function cancel(uint256 streamIndexSeed) external instrument("cancel") useFuzzedStreamSender(streamIndexSeed) {
Expand All @@ -106,7 +106,7 @@ abstract contract LockupHandler is BaseHandler {
// Record the returned amount by adding it to the ghost variable `returnedAmountsSum`. This is needed to
// check invariants against the contract's balance.
uint128 returnedAmount = lockup.returnableAmountOf(currentStreamId);
store.addReturnedAmount(returnedAmount);
_storage.addReturnedAmount(returnedAmount);

// Cancel the stream.
lockup.cancel(currentStreamId);
Expand Down Expand Up @@ -210,6 +210,6 @@ abstract contract LockupHandler is BaseHandler {
lockup.transferFrom({ from: currentRecipient, to: newRecipient, tokenId: currentStreamId });

// Update the recipient associated with this stream id.
store.updateRecipient(currentStreamId, newRecipient);
_storage.updateRecipient(currentStreamId, newRecipient);
}
}
8 changes: 2 additions & 6 deletions ...variant/handlers/LockupHandlerStore.t.sol → ...riant/handlers/LockupHandlerStorage.t.sol
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
// SPDX-License-Identifier: UNLICENSED
pragma solidity >=0.8.13 <0.9.0;

/// @title LockupHandlerStore
/// @title LockupHandlerStorage
/// @dev Storage contract for the lockup handler streams.
contract LockupHandlerStore {
contract LockupHandlerStorage {
/*//////////////////////////////////////////////////////////////////////////
PUBLIC TEST VARIABLES
//////////////////////////////////////////////////////////////////////////*/
Expand Down Expand Up @@ -35,8 +35,4 @@ contract LockupHandlerStore {
function updateRecipient(uint256 streamId, address newRecipient) external {
streamIdsToRecipients[streamId] = newRecipient;
}

function updateSender(uint256 streamId, address newSender) external {
streamIdsToSenders[streamId] = newSender;
}
}
33 changes: 11 additions & 22 deletions test/invariant/handlers/LockupLinearCreateHandler.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ import { ISablierV2LockupLinear } from "src/interfaces/ISablierV2LockupLinear.so
import { Broker, LockupLinear } from "src/types/DataTypes.sol";

import { BaseHandler } from "./BaseHandler.t.sol";
import { LockupHandlerStore } from "./LockupHandlerStore.t.sol";
import { LockupHandlerStorage } from "./LockupHandlerStorage.t.sol";

/// @title LockupLinearCreateHandler
/// @dev This contract is a complement of {LockupLinearHandler}. The goal is to bias the invariant calls
Expand All @@ -20,21 +20,20 @@ contract LockupLinearCreateHandler is BaseHandler {
uint256 internal constant MAX_STREAM_COUNT = 100;

/*//////////////////////////////////////////////////////////////////////////
PUBLIC TEST CONTRACTS
TEST CONTRACTS
//////////////////////////////////////////////////////////////////////////*/

IERC20 public asset;
ISablierV2LockupLinear public linear;
LockupHandlerStore public store;
LockupHandlerStorage public _storage;

/*//////////////////////////////////////////////////////////////////////////
CONSTRUCTOR
//////////////////////////////////////////////////////////////////////////*/

constructor(IERC20 asset_, ISablierV2LockupLinear linear_, LockupHandlerStore store_) {
constructor(IERC20 asset_, ISablierV2LockupLinear linear_, LockupHandlerStorage _storage_) {
asset = asset_;
linear = linear_;
store = store_;
_storage = _storage_;
}

/*//////////////////////////////////////////////////////////////////////////
Expand All @@ -60,10 +59,6 @@ contract LockupLinearCreateHandler is BaseHandler {
uint128 totalAmount;
}

struct CreateWithDurationsVars {
uint256 streamId;
}

function createWithDurations(
CreateWithDurationsParams memory params
) public instrument("createWithDurations") useNewSender(params.sender) {
Expand All @@ -73,7 +68,7 @@ contract LockupLinearCreateHandler is BaseHandler {
params.totalAmount = boundUint128(params.totalAmount, 1, 1_000_000_000e18);

// We don't want to fuzz more than a certain number of streams.
if (store.lastStreamId() >= MAX_STREAM_COUNT) {
if (_storage.lastStreamId() >= MAX_STREAM_COUNT) {
return;
}

Expand All @@ -89,8 +84,7 @@ contract LockupLinearCreateHandler is BaseHandler {
asset.approve({ spender: address(linear), amount: params.totalAmount });

// Create the stream.
CreateWithDurationsVars memory vars;
vars.streamId = linear.createWithDurations({
uint256 streamId = linear.createWithDurations({
sender: params.sender,
recipient: params.recipient,
totalAmount: params.totalAmount,
Expand All @@ -101,7 +95,7 @@ contract LockupLinearCreateHandler is BaseHandler {
});

// Store the stream id.
store.pushStreamId(vars.streamId, params.sender, params.recipient);
_storage.pushStreamId(streamId, params.sender, params.recipient);
}

struct CreateWithRangeParams {
Expand All @@ -113,10 +107,6 @@ contract LockupLinearCreateHandler is BaseHandler {
uint128 totalAmount;
}

struct CreateWithRangeVars {
uint256 streamId;
}

function createWithRange(
CreateWithRangeParams memory params
) public instrument("createWithRange") useNewSender(params.sender) {
Expand All @@ -127,7 +117,7 @@ contract LockupLinearCreateHandler is BaseHandler {
params.totalAmount = boundUint128(params.totalAmount, 1, 1_000_000_000e18);

// We don't want to fuzz more than a certain number of streams.
if (store.lastStreamId() >= MAX_STREAM_COUNT) {
if (_storage.lastStreamId() >= MAX_STREAM_COUNT) {
return;
}

Expand All @@ -143,8 +133,7 @@ contract LockupLinearCreateHandler is BaseHandler {
asset.approve({ spender: address(linear), amount: params.totalAmount });

// Create the stream.
CreateWithRangeVars memory vars;
vars.streamId = linear.createWithRange({
uint256 streamId = linear.createWithRange({
sender: params.sender,
recipient: params.recipient,
totalAmount: params.totalAmount,
Expand All @@ -155,6 +144,6 @@ contract LockupLinearCreateHandler is BaseHandler {
});

// Store the stream id.
store.pushStreamId(vars.streamId, params.sender, params.recipient);
_storage.pushStreamId(streamId, params.sender, params.recipient);
}
}
6 changes: 3 additions & 3 deletions test/invariant/handlers/LockupLinearHandler.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import { IERC20 } from "@openzeppelin/token/ERC20/IERC20.sol";
import { ISablierV2LockupLinear } from "src/interfaces/ISablierV2LockupLinear.sol";

import { LockupHandler } from "./LockupHandler.t.sol";
import { LockupHandlerStore } from "./LockupHandlerStore.t.sol";
import { LockupHandlerStorage } from "./LockupHandlerStorage.t.sol";

/// @title LockupLinearHandler
/// @dev This contract and not {SablierV2LockupLinear} is exposed to Foundry for invariant testing. The point is
Expand All @@ -15,6 +15,6 @@ contract LockupLinearHandler is LockupHandler {
constructor(
IERC20 asset_,
ISablierV2LockupLinear linear_,
LockupHandlerStore store_
) LockupHandler(asset_, linear_, store_) {}
LockupHandlerStorage _storage_
) LockupHandler(asset_, linear_, _storage_) {}
}
Loading

0 comments on commit fa5c007

Please sign in to comment.