test: initial structure for covering all key scenarios

src/AuraLockerModule.sol

@@ -34,13 +34,16 @@ contract AuraLockerModule is
                                        ERRORS
     //////////////////////////////////////////////////////////////////////////*/
     error NotKeeper(address agent);
+    error NotGovernance(address agent);
 
     error ZeroAddressValue();
 
     error ModuleNotEnabled();
 
     error TxFromModuleFailed();
 
+    error NothingToLock(uint256 timestamp);
+
     /*//////////////////////////////////////////////////////////////////////////
                                        EVENTS
     //////////////////////////////////////////////////////////////////////////*/
@@ -61,13 +64,19 @@ contract AuraLockerModule is
         _;
     }
 
+    /// @notice Enforce that the function is called by governance only
+    modifier onlyGovernance() {
+        if (msg.sender != BALANCER_MULTISIG) revert NotGovernance(msg.sender);
+        _;
+    }
+
     /*//////////////////////////////////////////////////////////////////////////
                                   EXTERNAL METHODS
     //////////////////////////////////////////////////////////////////////////*/
 
     /// @notice Assigns a new keeper address
     /// @param _keeper The address of the new keeper
-    function setKeeper(address _keeper) external {
+    function setKeeper(address _keeper) external onlyGovernance {
         if (_keeper == address(0)) revert ZeroAddressValue();
 
         address oldKeeper = keeper;
@@ -85,8 +94,7 @@ contract AuraLockerModule is
         override
         returns (bool requiresLocking, bytes memory execPayload)
     {
-        // if (!SAFE.isModuleEnabled(address(this)))
-        //   return (false, bytes("AuraLocker module is not enabled"));
+        if (!_isModuleEnabled()) return (false, bytes("AuraLocker module is not enabled"));
 
         (, uint256 unlockable,,) = AURA_LOCKER.lockedBalances(address(SAFE));
 
@@ -99,7 +107,10 @@ contract AuraLockerModule is
 
     /// @notice The actual execution of the action determined by the `checkUpkeep` method (AURA locking)
     function performUpkeep(bytes calldata /* _performData */ ) external override onlyKeeper {
-        // if (!SAFE.isModuleEnabled(address(this))) revert ModuleNotEnabled();
+        if (!_isModuleEnabled()) revert ModuleNotEnabled();
+
+        (, uint256 unlockable,,) = AURA_LOCKER.lockedBalances(address(SAFE));
+        if (unlockable == 0) revert NothingToLock(block.timestamp);
 
         // execute: `processExpiredLocks` via module
         if (
@@ -108,4 +119,13 @@ contract AuraLockerModule is
             )
         ) revert TxFromModuleFailed();
     }
+
+    /// @dev The Gnosis Safe version 1.1.1 does not expose directly `isModuleEnabled` method, so we need a workaround
+    function _isModuleEnabled() internal view returns (bool) {
+        address[] memory modules = SAFE.getModules();
+        for (uint256 i = 0; i < modules.length; i++) {
+            if (modules[i] == address(this)) return true;
+        }
+        return false;
+    }
 }

test/AuraLockerModule.t.sol

@@ -8,7 +8,7 @@ import {ILockAura} from "../src/interfaces/aura/ILockAura.sol";
 import {AuraLockerModule} from "../src/AuraLockerModule.sol";
 
 contract AuraLockerModuleTest is BaseFixture {
-    function test_checkUpkeep_when_NotNewUnlock() public {
+    function test_checkUpkeep_when_NotNewUnlock() public view {
         // at present nothing to lock
         (bool requiresLocking, bytes memory execPayload) = auraLockerModule.checkUpkeep(bytes(""));
         assertFalse(requiresLocking);
@@ -27,16 +27,58 @@ contract AuraLockerModuleTest is BaseFixture {
         // `disableModule(address prevModule, address module)`
         vm.prank(address(SAFE));
         SAFE.disableModule(address(1), address(auraLockerModule));
-        // assertFalse(SAFE.isModuleEnabled(address(auraLockerModule)));
+        address[] memory modules = SAFE.getModules();
+        for (uint256 i = 0; i < modules.length; i++) {
+            if (modules[i] == address(auraLockerModule)) assertFalse(true);
+        }
+
+        // once module its removed, the keeper trying to call `performUpkeep` should revert
+        vm.prank(auraLockerModule.keeper());
+        vm.expectRevert(abi.encodeWithSelector(AuraLockerModule.ModuleNotEnabled.selector));
+        auraLockerModule.performUpkeep(bytes(""));
     }
 
-    function test_revertWhen_NothingToUnlock() public {}
+    function testPerformUpkeep_revertWhen_NothingToLock() public {
+        // force a `performUpkeep` when weeks did not go by
+        skip(1 weeks);
+
+        vm.prank(auraLockerModule.keeper());
+        vm.expectRevert(abi.encodeWithSelector(AuraLockerModule.NothingToLock.selector, block.timestamp));
+        auraLockerModule.performUpkeep(bytes(""));
+    }
+
+    function testPerformUpkeepSuccess() public {
+        // move to future
+        skip(16 weeks);
+        (bool requiresLocking,) = auraLockerModule.checkUpkeep(bytes(""));
+        assertTrue(requiresLocking);
+
+        (uint256 totalAuraInLocker,, uint256 lockedBeforePerformUpkeep,) = AURA_LOCKER.lockedBalances(address(SAFE));
 
-    function testPerformUpkeepSuccess() public {}
+        vm.prank(auraLockerModule.keeper());
+        auraLockerModule.performUpkeep(bytes(""));
+
+        // check 2M where locked properly
+        (,, uint256 lockedAfterPerformUpkeep,) = AURA_LOCKER.lockedBalances(address(SAFE));
+        assertGt(lockedAfterPerformUpkeep, lockedBeforePerformUpkeep);
+        assertEq(totalAuraInLocker, lockedAfterPerformUpkeep);
+    }
 
     function testPerformUpkeep_revertWhen_NotKeeper() public {
         vm.prank(address(454545));
         vm.expectRevert(abi.encodeWithSelector(AuraLockerModule.NotKeeper.selector, address(454545)));
         auraLockerModule.performUpkeep(bytes(""));
     }
+
+    function testSetKeeper_revertWhen_NotGovernance() public {
+        vm.prank(address(454545));
+        vm.expectRevert(abi.encodeWithSelector(AuraLockerModule.NotGovernance.selector, address(454545)));
+        auraLockerModule.setKeeper(address(454545));
+    }
+
+    function testSetKeeper_revertWhen_AddressIsZero() public {
+        vm.prank(address(SAFE));
+        vm.expectRevert(abi.encodeWithSelector(AuraLockerModule.ZeroAddressValue.selector));
+        auraLockerModule.setKeeper(address(0));
+    }
 }

test/BaseFixture.sol

@@ -60,10 +60,11 @@ contract BaseFixture is Test {
         uint256 upkeepId = CL_REGISTRAR.registerUpkeep(registrationParams);
         address keeper = CL_REGISTRY.getForwarder(upkeepId);
         assertNotEq(keeper, address(0));
+        vm.stopPrank();
 
+        vm.prank(address(SAFE));
         auraLockerModule.setKeeper(keeper);
         assertEq(auraLockerModule.keeper(), keeper);
-        vm.stopPrank();
 
         _labelKeyContracts();
     }