| 1 | // SPDX-License-Identifier: MIT |
| 2 | pragma solidity >=0.7.0 <0.9.0; |
| 3 | |
| 4 | import {StdAssertions} from "../src/StdAssertions.sol"; |
| 5 | import {Vm} from "../src/Vm.sol"; |
| 6 | |
| 7 | interface VmInternal is Vm { |
| 8 | function _expectCheatcodeRevert(bytes memory message) external; |
| 9 | } |
| 10 | |
| 11 | contract StdAssertionsTest is StdAssertions { |
| 12 | string constant errorMessage = "User provided message"; |
| 13 | uint256 constant maxDecimals = 77; |
| 14 | |
| 15 | bool constant SHOULD_REVERT = true; |
| 16 | bool constant SHOULD_RETURN = false; |
| 17 | |
| 18 | bool constant STRICT_REVERT_DATA = true; |
| 19 | bool constant NON_STRICT_REVERT_DATA = false; |
| 20 | |
| 21 | VmInternal constant vm = VmInternal(address(uint160(uint256(keccak256("hevm cheat code"))))); |
| 22 | |
| 23 | function testFuzz_AssertEqCall_Return_Pass( |
| 24 | bytes memory callDataA, |
| 25 | bytes memory callDataB, |
| 26 | bytes memory returnData, |
| 27 | bool strictRevertData |
| 28 | ) external { |
| 29 | address targetA = address(new TestMockCall(returnData, SHOULD_RETURN)); |
| 30 | address targetB = address(new TestMockCall(returnData, SHOULD_RETURN)); |
| 31 | |
| 32 | assertEqCall(targetA, callDataA, targetB, callDataB, strictRevertData); |
| 33 | } |
| 34 | |
| 35 | function testFuzz_RevertWhenCalled_AssertEqCall_Return_Fail( |
| 36 | bytes memory callDataA, |
| 37 | bytes memory callDataB, |
| 38 | bytes memory returnDataA, |
| 39 | bytes memory returnDataB, |
| 40 | bool strictRevertData |
| 41 | ) external { |
| 42 | vm.assume(keccak256(returnDataA) != keccak256(returnDataB)); |
| 43 | |
| 44 | address targetA = address(new TestMockCall(returnDataA, SHOULD_RETURN)); |
| 45 | address targetB = address(new TestMockCall(returnDataB, SHOULD_RETURN)); |
| 46 | |
| 47 | vm._expectCheatcodeRevert( |
| 48 | bytes( |
| 49 | string.concat( |
| 50 | "Call return data does not match: ", vm.toString(returnDataA), " != ", vm.toString(returnDataB) |
| 51 | ) |
| 52 | ) |
| 53 | ); |
| 54 | assertEqCall(targetA, callDataA, targetB, callDataB, strictRevertData); |
| 55 | } |
| 56 | |
| 57 | function testFuzz_AssertEqCall_Revert_Pass( |
| 58 | bytes memory callDataA, |
| 59 | bytes memory callDataB, |
| 60 | bytes memory revertDataA, |
| 61 | bytes memory revertDataB |
| 62 | ) external { |
| 63 | address targetA = address(new TestMockCall(revertDataA, SHOULD_REVERT)); |
| 64 | address targetB = address(new TestMockCall(revertDataB, SHOULD_REVERT)); |
| 65 | |
| 66 | assertEqCall(targetA, callDataA, targetB, callDataB, NON_STRICT_REVERT_DATA); |
| 67 | } |
| 68 | |
| 69 | function testFuzz_RevertWhenCalled_AssertEqCall_Revert_Fail( |
| 70 | bytes memory callDataA, |
| 71 | bytes memory callDataB, |
| 72 | bytes memory revertDataA, |
| 73 | bytes memory revertDataB |
| 74 | ) external { |
| 75 | vm.assume(keccak256(revertDataA) != keccak256(revertDataB)); |
| 76 | |
| 77 | address targetA = address(new TestMockCall(revertDataA, SHOULD_REVERT)); |
| 78 | address targetB = address(new TestMockCall(revertDataB, SHOULD_REVERT)); |
| 79 | |
| 80 | vm._expectCheatcodeRevert( |
| 81 | bytes( |
| 82 | string.concat( |
| 83 | "Call revert data does not match: ", vm.toString(revertDataA), " != ", vm.toString(revertDataB) |
| 84 | ) |
| 85 | ) |
| 86 | ); |
| 87 | assertEqCall(targetA, callDataA, targetB, callDataB, STRICT_REVERT_DATA); |
| 88 | } |
| 89 | |
| 90 | function testFuzz_RevertWhenCalled_AssertEqCall_Fail( |
| 91 | bytes memory callDataA, |
| 92 | bytes memory callDataB, |
| 93 | bytes memory returnDataA, |
| 94 | bytes memory returnDataB, |
| 95 | bool strictRevertData |
| 96 | ) external { |
| 97 | address targetA = address(new TestMockCall(returnDataA, SHOULD_RETURN)); |
| 98 | address targetB = address(new TestMockCall(returnDataB, SHOULD_REVERT)); |
| 99 | |
| 100 | vm.expectRevert(bytes("assertion failed")); |
| 101 | this.assertEqCallExternal(targetA, callDataA, targetB, callDataB, strictRevertData); |
| 102 | |
| 103 | vm.expectRevert(bytes("assertion failed")); |
| 104 | this.assertEqCallExternal(targetB, callDataB, targetA, callDataA, strictRevertData); |
| 105 | } |
| 106 | |
| 107 | // Helper function to test outcome of assertEqCall via `expect` cheatcodes |
| 108 | function assertEqCallExternal( |
| 109 | address targetA, |
| 110 | bytes memory callDataA, |
| 111 | address targetB, |
| 112 | bytes memory callDataB, |
| 113 | bool strictRevertData |
| 114 | ) public { |
| 115 | assertEqCall(targetA, callDataA, targetB, callDataB, strictRevertData); |
| 116 | } |
| 117 | } |
| 118 | |
| 119 | contract TestMockCall { |
| 120 | bytes returnData; |
| 121 | bool shouldRevert; |
| 122 | |
| 123 | constructor(bytes memory returnData_, bool shouldRevert_) { |
| 124 | returnData = returnData_; |
| 125 | shouldRevert = shouldRevert_; |
| 126 | } |
| 127 | |
| 128 | fallback() external payable { |
| 129 | bytes memory returnData_ = returnData; |
| 130 | |
| 131 | if (shouldRevert) { |
| 132 | assembly { |
| 133 | revert(add(returnData_, 0x20), mload(returnData_)) |
| 134 | } |
| 135 | } else { |
| 136 | assembly { |
| 137 | return(add(returnData_, 0x20), mload(returnData_)) |
| 138 | } |
| 139 | } |
| 140 | } |
| 141 | } |