Merge pull request #141 from rhinestonewtf/feature/expect-revert-reason

Feature/expect revert reason

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@rhinestone/modulekit",
-  "version": "0.4.10",
+  "version": "0.4.11",
   "description": "A development kit for building and testing smart account modules.",
   "license": "GPL-3.0",
   "author": {

src/test/ModuleKitHelpers.sol

@@ -170,7 +170,6 @@ library ModuleKitHelpers {
         address module
     )
         internal
-        view
         returns (bool)
     {
         return HelperBase(instance.accountHelper).isModuleInstalled(instance, moduleTypeId, module);
@@ -183,7 +182,6 @@ library ModuleKitHelpers {
         bytes memory data
     )
         internal
-        view
         returns (bool)
     {
         return HelperBase(instance.accountHelper).isModuleInstalled(
@@ -312,7 +310,15 @@ library ModuleKitHelpers {
     //////////////////////////////////////////////////////////////////////////*/
 
     function expect4337Revert(AccountInstance memory) internal {
-        writeExpectRevert(1);
+        writeExpectRevert("");
+    }
+
+    function expect4337Revert(AccountInstance memory, bytes4 selector) internal {
+        writeExpectRevert(abi.encodePacked(selector));
+    }
+
+    function expect4337Revert(AccountInstance memory, bytes memory message) internal {
+        writeExpectRevert(message);
     }
 
     /**

src/test/helpers/HelperBase.sol

@@ -270,8 +270,8 @@ abstract contract HelperBase {
         address module
     )
         public
-        view
         virtual
+        deployAccountForAction(instance)
         returns (bool)
     {
         return isModuleInstalled(instance, moduleTypeId, module, "");
@@ -284,8 +284,8 @@ abstract contract HelperBase {
         bytes memory additionalContext
     )
         public
-        view
         virtual
+        deployAccountForAction(instance)
         returns (bool)
     {
         return IERC7579Account(instance.account).isModuleInstalled(

src/test/helpers/KernelHelpers.sol

@@ -394,9 +394,9 @@ contract KernelHelpers is HelperBase {
         bytes memory data
     )
         public
-        view
         virtual
         override
+        deployAccountForAction(instance)
         returns (bool)
     {
         if (moduleTypeId == MODULE_TYPE_HOOK) {

src/test/helpers/SafeHelpers.sol

@@ -245,9 +245,9 @@ contract SafeHelpers is HelperBase {
         bytes memory data
     )
         public
-        view
         virtual
         override
+        deployAccountForAction(instance)
         returns (bool)
     {
         if (moduleTypeId == MODULE_TYPE_HOOK) {

src/test/utils/ERC4337Helpers.sol

@@ -16,7 +16,8 @@ import { GasParser } from "./gas/GasParser.sol";
 import {
     getSimulateUserOp,
     getExpectRevert,
-    writeExpectRevert,
+    getExpectRevertMessage,
+    clearExpectRevert,
     getGasIdentifier,
     writeGasIdentifier,
     writeInstalledModule,
@@ -31,6 +32,8 @@ library ERC4337Helpers {
     error UserOperationReverted(
         bytes32 userOpHash, address sender, uint256 nonce, bytes revertReason
     );
+    error InvalidRevertMessage(bytes4 expected, bytes4 reason);
+    error InvalidRevertMessageBytes(bytes expected, bytes reason);
 
     function exec4337(PackedUserOperation[] memory userOps, IEntryPoint onEntryPoint) internal {
         uint256 isExpectRevert = getExpectRevert();
@@ -49,10 +52,12 @@ library ERC4337Helpers {
         // Execute userOps
         address payable beneficiary = payable(address(0x69));
         bytes memory userOpCalldata = abi.encodeCall(IEntryPoint.handleOps, (userOps, beneficiary));
-        (bool success,) = address(onEntryPoint).call(userOpCalldata);
+        (bool success, bytes memory returnData) = address(onEntryPoint).call(userOpCalldata);
 
         if (isExpectRevert == 0) {
             require(success, "UserOperation execution failed");
+        } else if (isExpectRevert == 2 && !success) {
+            checkRevertMessage(returnData);
         }
 
         // Parse logs and determine if a revert happened
@@ -68,14 +73,17 @@ library ERC4337Helpers {
                     abi.decode(logs[i].data, (uint256, bool, uint256, uint256));
                 totalUserOpGas = actualGasUsed;
                 if (!userOpSuccess) {
+                    bytes32 userOpHash = logs[i].topics[1];
                     if (isExpectRevert == 0) {
-                        bytes32 userOpHash = logs[i].topics[1];
                         bytes memory revertReason = getUserOpRevertReason(logs, userOpHash);
                         revert UserOperationReverted(
                             userOpHash, address(bytes20(logs[i].topics[2])), nonce, revertReason
                         );
                     } else {
-                        writeExpectRevert(0);
+                        if (isExpectRevert == 2) {
+                            checkRevertMessage(getUserOpRevertReason(logs, userOpHash));
+                        }
+                        clearExpectRevert();
                     }
                 }
             }
@@ -115,7 +123,7 @@ library ERC4337Helpers {
                 require(!success, "UserOperation execution did not fail as expected");
             }
         }
-        writeExpectRevert(0);
+        clearExpectRevert();
 
         // Calculate gas for userOp
         string memory gasIdentifier = getGasIdentifier();
@@ -139,7 +147,7 @@ library ERC4337Helpers {
 
     function getUserOpRevertReason(
         VmSafe.Log[] memory logs,
-        bytes32 /* userOpHash */
+        bytes32 userOpHash
     )
         internal
         pure
@@ -150,12 +158,29 @@ library ERC4337Helpers {
             if (
                 logs[i].topics[0]
                     == 0x1c4fada7374c0a9ee8841fc38afe82932dc0f8e69012e927f061a8bae611a201
+                    && logs[i].topics[1] == userOpHash
             ) {
                 (, revertReason) = abi.decode(logs[i].data, (uint256, bytes));
             }
         }
     }
 
+    function checkRevertMessage(bytes memory actualReason) internal view {
+        bytes memory revertMessage = getExpectRevertMessage();
+
+        if (revertMessage.length == 4) {
+            bytes4 expected = bytes4(revertMessage);
+            bytes4 actual = bytes4(actualReason);
+            if (expected != actual) {
+                revert InvalidRevertMessage(expected, actual);
+            }
+        } else {
+            if (revertMessage.length != actualReason.length) {
+                revert InvalidRevertMessageBytes(revertMessage, actualReason);
+            }
+        }
+    }
+
     function calculateGas(
         PackedUserOperation[] memory userOps,
         IEntryPoint onEntryPoint,

src/test/utils/Storage.sol

@@ -5,8 +5,18 @@ pragma solidity ^0.8.23;
                         EXPECT REVERT
 //////////////////////////////////////////////////////////////*/
 
-function writeExpectRevert(uint256 value) {
-    bytes32 slot = keccak256("ModuleKit.ExpectSlot");
+function writeExpectRevert(bytes memory message) {
+    uint256 value = 1;
+    bytes32 slot = keccak256("ModuleKit.ExpectMessageSlot");
+
+    if (message.length > 0) {
+        value = 2;
+        assembly {
+            sstore(slot, message)
+        }
+    }
+
+    slot = keccak256("ModuleKit.ExpectSlot");
     assembly {
         sstore(slot, value)
     }
@@ -19,6 +29,25 @@ function getExpectRevert() view returns (uint256 value) {
     }
 }
 
+function getExpectRevertMessage() view returns (bytes memory data) {
+    bytes32 slot = keccak256("ModuleKit.ExpectMessageSlot");
+    assembly {
+        data := sload(slot)
+    }
+}
+
+function clearExpectRevert() {
+    bytes32 slot = keccak256("ModuleKit.ExpectSlot");
+    assembly {
+        sstore(slot, 0)
+    }
+
+    slot = keccak256("ModuleKit.ExpectMessageSlot");
+    assembly {
+        sstore(slot, 0)
+    }
+}
+
 /*//////////////////////////////////////////////////////////////
                         GAS IDENTIFIER
 //////////////////////////////////////////////////////////////*/