feat: add script to upgrade Pool Voter

contracts/interfaces/IOmnichainStaking.sol

@@ -19,7 +19,7 @@ interface IOmnichainStaking is IVotes {
     // An omni-chain staking contract that allows users to stake their veNFT
     // and get some voting power. Once staked the voting power is available cross-chain.
 
-    function unstakeToken(uint256 tokenId) external;
+    function unstakeAndWithdraw(uint256 tokenId) external;
 
     // function totalSupply() external view returns (uint256);
 

contracts/locker/BaseLocker.sol

@@ -292,7 +292,8 @@ contract BaseLocker is
     /// @dev Only possible if the lock has expired
     function withdraw(uint256 _tokenId) public virtual nonReentrant {
         require(
-            _isAuthorized(ownerOf(_tokenId), msg.sender, _tokenId),
+            _isAuthorized(ownerOf(_tokenId), msg.sender, _tokenId) ||
+            address(staking) == msg.sender,
             "caller is not owner nor approved"
         );
 

contracts/locker/staking/OmnichainStakingBase.sol

@@ -165,27 +165,13 @@ abstract contract OmnichainStakingBase is
     }
 
     /**
-     * @dev Unstakes a regular token NFT and transfers it back to the user.
-     * @param tokenId The ID of the regular token NFT to unstake.
+     * @notice A single withdraw function to unstake NFT, transfer it back to
+     * the user, and also withdraw all tokens for `tokenId` from the locker.
+     * @param tokenId The ID of the regular token NFT to unstake and withdraw.
      */
-    function unstakeToken(uint256 tokenId) external updateReward(msg.sender) {
-        require(lockedByToken[tokenId] != address(0), "!tokenId");
-        address lockedBy_ = lockedByToken[tokenId];
-        if (_msgSender() != lockedBy_)
-            revert InvalidUnstaker(_msgSender(), lockedBy_);
-
-        delete lockedByToken[tokenId];
-        lockedTokenIdNfts[_msgSender()] = deleteAnElement(
-            lockedTokenIdNfts[_msgSender()],
-            tokenId
-        );
-
-        // reset and burn voting power
-        _burn(msg.sender, tokenPower[tokenId]);
-        tokenPower[tokenId] = 0;
-        votingPowerCombined.reset(msg.sender);
-
-        locker.safeTransferFrom(address(this), msg.sender, tokenId);
+    function unstakeAndWithdraw(uint256 tokenId) external {
+        _unstakeToken(tokenId);
+        locker.withdraw(tokenId);
     }
 
     /**
@@ -408,6 +394,30 @@ abstract contract OmnichainStakingBase is
         }
     }
 
+    /**
+     * @dev Unstakes a regular token NFT and transfers it back to the user.
+     * @param tokenId The ID of the regular token NFT to unstake.
+     */
+    function _unstakeToken(uint256 tokenId) internal updateReward(msg.sender) {
+        require(lockedByToken[tokenId] != address(0), "!tokenId");
+        address lockedBy_ = lockedByToken[tokenId];
+        if (_msgSender() != lockedBy_)
+            revert InvalidUnstaker(_msgSender(), lockedBy_);
+
+        delete lockedByToken[tokenId];
+        lockedTokenIdNfts[_msgSender()] = deleteAnElement(
+            lockedTokenIdNfts[_msgSender()],
+            tokenId
+        );
+
+        // reset and burn voting power
+        _burn(msg.sender, tokenPower[tokenId]);
+        tokenPower[tokenId] = 0;
+        votingPowerCombined.reset(msg.sender);
+
+        locker.safeTransferFrom(address(this), msg.sender, tokenId);
+    }
+
     /**
      * @dev Deletes an element from an array.
      * @param elements The array to delete from.

contracts/voter/PoolVoter.sol

@@ -42,11 +42,12 @@ contract PoolVoter is
      * @param _staking The address of the staking token (VE token).
      * @param _reward The address of the reward token.
      */
-    function init(address _staking, address _reward) external reinitializer(1) {
+    function init(address _staking, address _reward, address _votingPowerCombined) external reinitializer(1) {
         staking = IVotes(_staking);
         reward = IERC20(_reward);
         __ReentrancyGuard_init();
         __Ownable_init(msg.sender);
+        votingPowerCombined = _votingPowerCombined;
     }
 
     /**

scripts/upgrade-poolVoter.ts

@@ -0,0 +1,28 @@
+import {ethers, upgrades} from "hardhat";
+
+const ZERO_TOKEN_ADDRESS = "0x78354f8DcCB269a615A7e0a24f9B0718FDC3C7A7";
+const OMNICHAIN_STAKING_ADDRESS = "0xf374229a18ff691406f99CCBD93e8a3f16B68888";
+
+async function main() {
+
+  const poolVoterProxy = "0x5346e9ab27D7874Db95993667D1Cb8338913f0aF"
+  const poolVoter = await ethers.getContractFactory("PoolVoter");
+
+  console.log("Upgradig to new PoolVoter implementation");
+
+  if (ZERO_TOKEN_ADDRESS.length && OMNICHAIN_STAKING_ADDRESS.length) {
+
+    const upgradedPoolVoter = await upgrades.upgradeProxy(poolVoterProxy, poolVoter);
+    const tx = await upgradedPoolVoter.init(OMNICHAIN_STAKING_ADDRESS, ZERO_TOKEN_ADDRESS);
+    await  tx.wait();
+
+    console.log("PoolVoter upgraded to", upgradedPoolVoter.address);
+  } else {
+    throw new Error("Invalid init arguments");
+  }
+}
+
+main().catch((error) => {
+    console.error(error);
+    process.exitCode = 1;
+});

test/fork/VotingPowerCombined/linea.test.ts

@@ -0,0 +1,54 @@
+import { expect, should } from "chai";
+import { e18, initMainnetUser } from "../../fixtures/utils";
+import { VestedZeroNFT, VotingPowerCombined } from "../../../typechain-types";
+import { SignerWithAddress } from "@nomicfoundation/hardhat-ethers/signers";
+import { time } from "@nomicfoundation/hardhat-network-helpers";
+import { setForkBlock } from "../utils";
+import { getNetworkDetails } from "../constants";
+import {
+  Contract,
+  ContractTransactionResponse,
+  parseEther,
+  parseUnits,
+} from "ethers";
+import { ethers } from "hardhat";
+import { getGovernanceContracts } from "../helper";
+
+const FORK = process.env.FORK === "true";
+const FORKED_NETWORK = process.env.FORKED_NETWORK ?? "";
+const BLOCK_NUMBER = 5969983;
+
+const STAKING_ADDRESS = "0xf374229a18ff691406f99CCBD93e8a3f16B68888";
+const REWARD_ADDRESS = "0x78354f8DcCB269a615A7e0a24f9B0718FDC3C7A7";
+const OWNER = "0x0F6e98A756A40dD050dC78959f45559F98d3289d";
+
+if (FORK) {
+  let votingPowerCombined: VotingPowerCombined;
+  let deployerForked: SignerWithAddress;
+  describe.only("VotingPowerCombined ForkTests", async () => {
+    beforeEach(async () => {
+      votingPowerCombined = await ethers.getContractAt(
+        "VotingPowerCombined",
+        "0x2666951A62d82860E8e1385581E2FB7669097647"
+      );
+      [deployerForked] = await ethers.getSigners();
+      await setForkBlock(BLOCK_NUMBER);
+
+      //
+
+      const poolVoterFactory = await ethers.getContractFactory("PoolVoter");
+      const poolVoter = await poolVoterFactory.deploy();
+      await poolVoter.init(STAKING_ADDRESS, REWARD_ADDRESS, votingPowerCombined.target);
+
+      const owner = await initMainnetUser(OWNER);
+      await votingPowerCombined.connect(owner).setAddresses(STAKING_ADDRESS, REWARD_ADDRESS, poolVoter.target);
+    });
+
+    it("Should reset voting power", async () => {
+      const resetterWallet = await initMainnetUser("0x7Ff4e6A2b7B43cEAB1fC07B0CBa00f834846ADEd", parseEther('100'));
+
+      const resetTransaction = votingPowerCombined.connect(resetterWallet).reset(resetterWallet.address);
+      await expect(resetTransaction).to.not.be.revertedWith('Invalid reset performed');
+    });
+  });
+}