diff --git a/src/L2/L2VotingPower.sol b/src/L2/L2VotingPower.sol index a8f358d4..2742bf37 100644 --- a/src/L2/L2VotingPower.sol +++ b/src/L2/L2VotingPower.sol @@ -79,15 +79,25 @@ contract L2VotingPower is ERC20VotesUpgradeable, OwnableUpgradeable, UUPSUpgrade /// @param positionAfter Locking position after the adjustment. function adjustVotingPower( address ownerAddress, - LockingPosition memory positionBefore, - LockingPosition memory positionAfter + LockingPosition[] memory positionBefore, + LockingPosition[] memory positionAfter ) public virtual onlyStakingContract { - _mint(ownerAddress, votingPower(positionAfter)); - _burn(ownerAddress, votingPower(positionBefore)); + require( + positionBefore.length == 0 || positionBefore.length == 1, "L2VotingPower: invalid positionBefore length" + ); + require(positionAfter.length == 0 || positionAfter.length == 1, "L2VotingPower: invalid positionAfter length"); + + if (positionAfter.length > 0) { + _mint(ownerAddress, votingPower(positionAfter[0])); + } + + if (positionBefore.length > 0) { + _burn(ownerAddress, votingPower(positionBefore[0])); + } } /// @notice Overrides clock() function to make the token & governor timestamp-based diff --git a/test/L2/L2VotingPower.t.sol b/test/L2/L2VotingPower.t.sol index dccbd070..7c9f2963 100644 --- a/test/L2/L2VotingPower.t.sol +++ b/test/L2/L2VotingPower.t.sol @@ -33,6 +33,8 @@ contract L2VotingPowerTest is Test { L2VotingPower public l2VotingPower; address stakingContractAddress; + LockingPosition[] public positionBefore; + LockingPosition[] public positionAfter; function setUp() public { utils = new Utils(); @@ -80,8 +82,8 @@ contract L2VotingPowerTest is Test { } function test_AdjustVotingPower_EventsAreEmitted() public { - LockingPosition memory positionBefore = LockingPosition(50, 50, 50); - LockingPosition memory positionAfter = LockingPosition(100, 100, 100); + positionBefore.push(LockingPosition(50, 50, 50)); + positionAfter.push(LockingPosition(100, 100, 100)); // check that events for mint and burn are emitted vm.expectEmit(true, true, true, true); @@ -95,14 +97,71 @@ contract L2VotingPowerTest is Test { } function test_AdjustVotingPower_NotStakingContract() public { - LockingPosition memory positionBefore = LockingPosition(50, 50, 50); - LockingPosition memory positionAfter = LockingPosition(100, 100, 100); + positionBefore.push(LockingPosition(50, 50, 50)); + positionAfter.push(LockingPosition(100, 100, 100)); // call it as non-staking contract vm.expectRevert("L2VotingPower: only staking contract can call this function"); l2VotingPower.adjustVotingPower(address(this), positionBefore, positionAfter); } + function test_AdjustVotingPower_PositionBeforeLength() public { + positionBefore.push(LockingPosition(50, 50, 50)); + positionBefore.push(LockingPosition(50, 50, 50)); + positionAfter.push(LockingPosition(100, 100, 100)); + + // call it as staking contract + vm.prank(stakingContractAddress); + vm.expectRevert("L2VotingPower: invalid positionBefore length"); + l2VotingPower.adjustVotingPower(address(this), positionBefore, positionAfter); + } + + function test_AdjustVotingPower_PositionAfterLength() public { + positionBefore.push(LockingPosition(50, 50, 50)); + positionAfter.push(LockingPosition(100, 100, 100)); + positionAfter.push(LockingPosition(100, 100, 100)); + + // call it as staking contract + vm.prank(stakingContractAddress); + vm.expectRevert("L2VotingPower: invalid positionAfter length"); + l2VotingPower.adjustVotingPower(address(this), positionBefore, positionAfter); + } + + function test_AdjustVotingPower_PositionBeforeIsNull() public { + positionAfter.push(LockingPosition(100, 100, 100)); + + // check that event for mint is emitted + vm.expectEmit(true, true, true, true); + emit IERC20.Transfer(address(0), address(this), 6250); + + // call it as staking contract + vm.prank(stakingContractAddress); + l2VotingPower.adjustVotingPower(address(this), positionBefore, positionAfter); + } + + function test_AdjustVotingPower_PositionAfterIsNull() public { + // mint some tokens that then can be burned + positionAfter.push(LockingPosition(500, 500, 500)); + vm.prank(stakingContractAddress); + l2VotingPower.adjustVotingPower(address(this), positionBefore, positionAfter); + + // clear positionAfter array + positionAfter.pop(); + assertEq(positionAfter.length, 0); + assertEq(positionBefore.length, 0); + + // only positionBefore is set + positionBefore.push(LockingPosition(50, 50, 50)); + + // check that event for burn is emitted + vm.expectEmit(true, true, true, true); + emit IERC20.Transfer(address(this), address(0), 2500); + + // call it as staking contract + vm.prank(stakingContractAddress); + l2VotingPower.adjustVotingPower(address(this), positionBefore, positionAfter); + } + function test_Clock() public { assertEq(l2VotingPower.clock(), vm.getBlockNumber());