fix: add tests

bgd-labs · Jul 17, 2024 · ac6984f · ac6984f
1 parent 20fd662
commit ac6984f
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 3 deletions.
diff --git a/src/contracts/access-control/UpgradableOwnableWithGuardian.sol b/src/contracts/access-control/UpgradableOwnableWithGuardian.sol
@@ -24,6 +24,16 @@ abstract contract UpgradableOwnableWithGuardian is OwnableUpgradeable, IWithGuar
     }
   }
 
+  /**
+   * @dev The caller account is not authorized to perform an operation.
+   */
+  error OnlyGuardianInvalidCaller(address account);
+
+  /**
+   * @dev The caller account is not authorized to perform an operation.
+   */
+  error OnlyGuardianOrOwnerInvalidCaller(address account);
+
   /**
    * @dev Initializes the contract setting the address provided by the deployer as the initial owner.
    */
@@ -63,10 +73,11 @@ abstract contract UpgradableOwnableWithGuardian is OwnableUpgradeable, IWithGuar
   }
 
   function _checkGuardian() internal view {
-    require(guardian() == _msgSender(), 'ONLY_BY_GUARDIAN');
+    if (guardian() != _msgSender()) revert OnlyGuardianInvalidCaller(_msgSender());
   }
 
   function _checkOwnerOrGuardian() internal view {
-    require(_msgSender() == owner() || _msgSender() == guardian(), 'ONLY_BY_OWNER_OR_GUARDIAN');
+    if (_msgSender() != owner() && _msgSender() != guardian())
+      revert OnlyGuardianOrOwnerInvalidCaller(_msgSender());
   }
 }
diff --git a/test/UpgradableOwnableWithGuardian.t copy.sol b/test/UpgradableOwnableWithGuardian.t copy.sol
@@ -32,10 +32,22 @@ contract TestOfUpgradableOwnableWithGuardian is Test {
   }
 
   function test_onlyGuardian() external {
+    vm.expectRevert(
+      abi.encodeWithSelector(
+        UpgradableOwnableWithGuardian.OnlyGuardianInvalidCaller.selector,
+        address(this)
+      )
+    );
     ImplOwnableWithGuardian(address(withGuardian)).mock_onlyGuardian();
   }
 
   function test_onlyOwnerOrGuardian() external {
+    vm.expectRevert(
+      abi.encodeWithSelector(
+        UpgradableOwnableWithGuardian.OnlyGuardianOrOwnerInvalidCaller.selector,
+        address(this)
+      )
+    );
     ImplOwnableWithGuardian(address(withGuardian)).mock_onlyOwnerOrGuardian();
   }
 
@@ -53,7 +65,12 @@ contract TestOfUpgradableOwnableWithGuardian is Test {
     vm.assume(newGuardian != owner && newGuardian != guardian);
 
     vm.prank(newGuardian);
-    vm.expectRevert('ONLY_BY_OWNER_OR_GUARDIAN');
+    vm.expectRevert(
+      abi.encodeWithSelector(
+        UpgradableOwnableWithGuardian.OnlyGuardianOrOwnerInvalidCaller.selector,
+        address(this)
+      )
+    );
     withGuardian.updateGuardian(newGuardian);
   }
 }