Let the staking provider to refresh KEEP stake owner

contracts/staking/IStaking.sol

@@ -70,7 +70,8 @@ interface IStaking {
         address authorizer
     ) external;
 
-    /// @notice Refresh Keep stake owner. Can be called only by the old owner.
+    /// @notice Refresh Keep stake owner. Can be called only by the old owner
+    ///         or their staking provider.
     /// @dev The staking provider in T staking contract is the legacy KEEP
     ///      staking contract operator.
     function refreshKeepStakeOwner(address stakingProvider) external;

contracts/staking/KeepStake.sol

@@ -28,12 +28,13 @@ interface IManagedGrant {
 ///         KEEP stakers to use their stakes in T network and weights them based
 ///         on KEEP<>T token ratio. KEEP stake owner is cached in T staking
 ///         contract and used to restrict access to all functions only owner or
-///         operator should call. To cache KEEP staking contract in T staking
-///         contract, it fitst needs to resolve the owner. Resolving liquid
-///         KEEP stake owner is easy. Resolving token grant stake owner is
-///         complicated and not possible to do on-chain from a contract external
-///         to KEEP TokenStaking contract. Keep TokenStaking knows the grant ID
-///         but does not expose it externally.
+///         operator should call. To cache KEEP stake owner in T staking
+///         contract, T staking contract first needs to resolve the owner.
+///
+///         Resolving liquid KEEP stake owner is easy. Resolving token grant
+///         stake owner is complicated and not possible to do on-chain from
+///         a contract external to KEEP TokenStaking contract. Keep TokenStaking
+///         knows the grant ID but does not expose it externally.
 ///
 ///         KeepStake contract addresses this problem by exposing
 ///         operator-owner mappings snapshotted off-chain based on events and

contracts/staking/TokenStaking.sol

@@ -397,17 +397,18 @@ contract TokenStaking is Initializable, IStaking, Checkpoints {
         );
     }
 
-    /// @notice Refresh Keep stake owner. Can be called only by the old owner.
+    /// @notice Refresh Keep stake owner. Can be called only by the old owner
+    ///         or their staking provider.
     /// @dev The staking provider in T staking contract is the legacy KEEP
     ///      staking contract operator.
-    function refreshKeepStakeOwner(address stakingProvider) external override {
+    function refreshKeepStakeOwner(address stakingProvider)
+        external
+        override
+        onlyOwnerOrStakingProvider(stakingProvider)
+    {
         StakingProviderInfo storage stakingProviderStruct = stakingProviders[
             stakingProvider
         ];
-        require(
-            stakingProviderStruct.owner == msg.sender,
-            "Caller is not owner"
-        );
         address newOwner = keepStake.resolveOwner(stakingProvider);
 
         emit OwnerRefreshed(

test/staking/TokenStaking.test.js

@@ -864,11 +864,11 @@ describe("TokenStaking", () => {
           tokenStaking
             .connect(stakingProvider)
             .refreshKeepStakeOwner(stakingProvider.address)
-        ).to.be.revertedWith("Caller is not owner")
+        ).to.be.revertedWith("Not owner or provider")
       })
     })
 
-    context("when caller is not old owner", () => {
+    context("when caller is neither old owner nor staking provider", () => {
       it("should revert", async () => {
         await tToken
           .connect(staker)
@@ -883,108 +883,122 @@ describe("TokenStaking", () => {
           )
         await expect(
           tokenStaking
-            .connect(stakingProvider)
+            .connect(authorizer)
             .refreshKeepStakeOwner(stakingProvider.address)
-        ).to.be.revertedWith("Caller is not owner")
+        ).to.be.revertedWith("Not owner or provider")
       })
     })
 
-    context("when grantee was not changed", () => {
-      let tx
+    const contextRefreshKeepStakeOwner = (getCaller) => {
+      context("when grantee was not changed", () => {
+        let tx
 
-      beforeEach(async () => {
-        const createdAt = 1
-        await keepStakingMock.setOperator(
-          stakingProvider.address,
-          staker.address,
-          beneficiary.address,
-          authorizer.address,
-          createdAt,
-          0,
-          initialStakerBalance
-        )
-        await keepStakingMock.setEligibility(
-          stakingProvider.address,
-          tokenStaking.address,
-          true
-        )
-        await tokenStaking.stakeKeep(stakingProvider.address)
+        beforeEach(async () => {
+          const createdAt = 1
+          await keepStakingMock.setOperator(
+            stakingProvider.address,
+            staker.address,
+            beneficiary.address,
+            authorizer.address,
+            createdAt,
+            0,
+            initialStakerBalance
+          )
+          await keepStakingMock.setEligibility(
+            stakingProvider.address,
+            tokenStaking.address,
+            true
+          )
+          await tokenStaking.stakeKeep(stakingProvider.address)
 
-        tx = await tokenStaking
-          .connect(staker)
-          .refreshKeepStakeOwner(stakingProvider.address)
-      })
+          tx = await tokenStaking
+            .connect(getCaller())
+            .refreshKeepStakeOwner(stakingProvider.address)
+        })
 
-      it("should not update owner", async () => {
-        expect(
-          await tokenStaking.rolesOf(stakingProvider.address)
-        ).to.deep.equal([
-          staker.address,
-          beneficiary.address,
-          authorizer.address,
-        ])
-      })
+        it("should not update owner", async () => {
+          expect(
+            await tokenStaking.rolesOf(stakingProvider.address)
+          ).to.deep.equal([
+            staker.address,
+            beneficiary.address,
+            authorizer.address,
+          ])
+        })
 
-      it("should emit OwnerRefreshed", async () => {
-        await expect(tx)
-          .to.emit(tokenStaking, "OwnerRefreshed")
-          .withArgs(stakingProvider.address, staker.address, staker.address)
+        it("should emit OwnerRefreshed", async () => {
+          await expect(tx)
+            .to.emit(tokenStaking, "OwnerRefreshed")
+            .withArgs(stakingProvider.address, staker.address, staker.address)
+        })
       })
-    })
 
-    context("when grantee was changed", () => {
-      let tx
+      context("when grantee was changed", () => {
+        let tx
 
-      beforeEach(async () => {
-        const createdAt = 1
-        await keepStakingMock.setOperator(
-          stakingProvider.address,
-          otherStaker.address,
-          beneficiary.address,
-          authorizer.address,
-          createdAt,
-          0,
-          initialStakerBalance
-        )
-        await keepStakingMock.setEligibility(
-          stakingProvider.address,
-          tokenStaking.address,
-          true
-        )
-        await tokenStaking.stakeKeep(stakingProvider.address)
+        beforeEach(async () => {
+          const createdAt = 1
+          await keepStakingMock.setOperator(
+            stakingProvider.address,
+            otherStaker.address,
+            beneficiary.address,
+            authorizer.address,
+            createdAt,
+            0,
+            initialStakerBalance
+          )
+          await keepStakingMock.setEligibility(
+            stakingProvider.address,
+            tokenStaking.address,
+            true
+          )
+          await tokenStaking.stakeKeep(stakingProvider.address)
 
-        await keepStakingMock.setOperator(
-          stakingProvider.address,
-          staker.address,
-          beneficiary.address,
-          authorizer.address,
-          createdAt,
-          0,
-          initialStakerBalance
-        )
-        tx = await tokenStaking
-          .connect(otherStaker)
-          .refreshKeepStakeOwner(stakingProvider.address)
+          await keepStakingMock.setOperator(
+            stakingProvider.address,
+            staker.address,
+            beneficiary.address,
+            authorizer.address,
+            createdAt,
+            0,
+            initialStakerBalance
+          )
+          tx = await tokenStaking
+            .connect(otherStaker)
+            .refreshKeepStakeOwner(stakingProvider.address)
+        })
+
+        it("should update owner", async () => {
+          expect(
+            await tokenStaking.rolesOf(stakingProvider.address)
+          ).to.deep.equal([
+            staker.address,
+            beneficiary.address,
+            authorizer.address,
+          ])
+        })
+
+        it("should emit OwnerRefreshed", async () => {
+          await expect(tx)
+            .to.emit(tokenStaking, "OwnerRefreshed")
+            .withArgs(
+              stakingProvider.address,
+              otherStaker.address,
+              staker.address
+            )
+        })
       })
+    }
 
-      it("should update owner", async () => {
-        expect(
-          await tokenStaking.rolesOf(stakingProvider.address)
-        ).to.deep.equal([
-          staker.address,
-          beneficiary.address,
-          authorizer.address,
-        ])
+    context("when caller is the old owner", () => {
+      contextRefreshKeepStakeOwner(() => {
+        return staker
       })
+    })
 
-      it("should emit OwnerRefreshed", async () => {
-        await expect(tx)
-          .to.emit(tokenStaking, "OwnerRefreshed")
-          .withArgs(
-            stakingProvider.address,
-            otherStaker.address,
-            staker.address
-          )
+    context("when caller is the staking provider", () => {
+      contextRefreshKeepStakeOwner(() => {
+        return stakingProvider
       })
     })
   })