Modify ThresholdStakingForTACoApplicationMock to be ownable and limit…

… changes to contract to only allow from owner.
derekpierre · Sep 25, 2023 · 33f3a41 · 33f3a41
1 parent 48c43e5
commit 33f3a41
Showing 1 changed file with 17 additions and 11 deletions.
diff --git a/contracts/test/TACoApplicationTestSet.sol b/contracts/test/TACoApplicationTestSet.sol
@@ -5,6 +5,7 @@ pragma solidity ^0.8.0;
 import "../contracts/TACoApplication.sol";
 import "@threshold/contracts/staking/IApplication.sol";
 import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+import "../.cache/openzeppelin/v4.9.1/access/Ownable.sol";
 
 /**
  * @notice Contract for testing the application contract
@@ -18,7 +19,7 @@ contract TToken is ERC20("T", "T") {
 /**
  * @notice Contract for testing TACo application contract
  */
-contract ThresholdStakingForTACoApplicationMock {
+contract ThresholdStakingForTACoApplicationMock is Ownable {
     struct StakingProviderInfo {
         address owner;
         address payable beneficiary;
@@ -36,7 +37,7 @@ contract ThresholdStakingForTACoApplicationMock {
     address public notifier;
     address[] public stakingProvidersToSeize;
 
-    function setApplication(IApplication _application) external {
+    function setApplication(IApplication _application) external onlyOwner {
         application = _application;
     }
 
@@ -49,7 +50,7 @@ contract ThresholdStakingForTACoApplicationMock {
         address _owner,
         address payable _beneficiary,
         address _authorizer
-    ) public {
+    ) public onlyOwner {
         StakingProviderInfo storage info = stakingProviderInfo[_stakingProvider];
         info.owner = _owner;
         info.beneficiary = _beneficiary;
@@ -60,15 +61,18 @@ contract ThresholdStakingForTACoApplicationMock {
      * @dev If the function is called with only the _stakingProvider parameter,
      * we presume that the caller wants that address set for the other roles as well.
      */
-    function setRoles(address _stakingProvider) external {
+    function setRoles(address _stakingProvider) external onlyOwner {
         setRoles(_stakingProvider, _stakingProvider, payable(_stakingProvider), _stakingProvider);
     }
 
-    function setAuthorized(address _stakingProvider, uint96 _authorized) external {
+    function setAuthorized(address _stakingProvider, uint96 _authorized) external onlyOwner {
         stakingProviderInfo[_stakingProvider].authorized = _authorized;
     }
 
-    function setDecreaseRequest(address _stakingProvider, uint96 _decreaseRequestTo) external {
+    function setDecreaseRequest(
+        address _stakingProvider,
+        uint96 _decreaseRequestTo
+    ) external onlyOwner {
         stakingProviderInfo[_stakingProvider].decreaseRequestTo = _decreaseRequestTo;
     }
 
@@ -90,7 +94,9 @@ contract ThresholdStakingForTACoApplicationMock {
         authorizer = info.authorizer;
     }
 
-    function approveAuthorizationDecrease(address _stakingProvider) external returns (uint96) {
+    function approveAuthorizationDecrease(
+        address _stakingProvider
+    ) external onlyOwner returns (uint96) {
         StakingProviderInfo storage info = stakingProviderInfo[_stakingProvider];
         info.authorized = info.decreaseRequestTo;
         return info.authorized;
@@ -101,7 +107,7 @@ contract ThresholdStakingForTACoApplicationMock {
         uint256 _rewardMultiplier,
         address _notifier,
         address[] memory _stakingProviders
-    ) external {
+    ) external onlyOwner {
         amountToSeize = _amount;
         rewardMultiplier = _rewardMultiplier;
         notifier = _notifier;
@@ -116,7 +122,7 @@ contract ThresholdStakingForTACoApplicationMock {
         address _stakingProvider,
         uint96 _fromAmount,
         uint96 _toAmount
-    ) external {
+    ) external onlyOwner {
         application.authorizationIncreased(_stakingProvider, _fromAmount, _toAmount);
         stakingProviderInfo[_stakingProvider].authorized = _toAmount;
     }
@@ -125,7 +131,7 @@ contract ThresholdStakingForTACoApplicationMock {
         address _stakingProvider,
         uint96 _fromAmount,
         uint96 _toAmount
-    ) external {
+    ) external onlyOwner {
         application.involuntaryAuthorizationDecrease(_stakingProvider, _fromAmount, _toAmount);
         stakingProviderInfo[_stakingProvider].authorized = _toAmount;
     }
@@ -134,7 +140,7 @@ contract ThresholdStakingForTACoApplicationMock {
         address _stakingProvider,
         uint96 _fromAmount,
         uint96 _toAmount
-    ) external {
+    ) external onlyOwner {
         application.authorizationDecreaseRequested(_stakingProvider, _fromAmount, _toAmount);
         stakingProviderInfo[_stakingProvider].decreaseRequestTo = _toAmount;
     }