From 775754f5f5d7d62f81dfe1ff95ba801bbb4eac39 Mon Sep 17 00:00:00 2001
From: baroooo <baranseltekin@gmail.com>
Date: Thu, 14 Nov 2024 13:58:20 +0100
Subject: [PATCH 01/16] feat: upgrade locking to work with celol2

---
 contracts/governance/locking/LockingBase.sol  | 103 ++++-
 test/fork/ForkTests.t.sol                     |   3 +
 test/unit/governance/Locking/LockingTest.sol  |   4 +
 .../governance/Locking/locking.upgrade.t.sol  | 384 ++++++++++++++++++
 test/utils/harnesses/LockingHarness.sol       |  16 +-
 5 files changed, 503 insertions(+), 7 deletions(-)
 create mode 100644 test/unit/governance/Locking/locking.upgrade.t.sol

diff --git a/contracts/governance/locking/LockingBase.sol b/contracts/governance/locking/LockingBase.sol
index f6e8509b..24211440 100644
--- a/contracts/governance/locking/LockingBase.sol
+++ b/contracts/governance/locking/LockingBase.sol
@@ -7,6 +7,8 @@ import "openzeppelin-contracts-upgradeable/contracts/access/OwnableUpgradeable.s
 import "openzeppelin-contracts-upgradeable/contracts/governance/utils/IVotesUpgradeable.sol";
 import "./libs/LibBrokenLine.sol";
 
+import "forge-std/console.sol";
+
 /**
  * @title LockingBase
  * @dev This abstract contract provides the foundational functionality
@@ -17,9 +19,15 @@ import "./libs/LibBrokenLine.sol";
 abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
   using LibBrokenLine for LibBrokenLine.BrokenLine;
   /**
-   * @dev Duration of a week in blocks on the CELO blockchain assuming 5 seconds per block
+   * @dev Duration of a week in blocks on the CELO blockchain before the L2 transition (5 seconds per block)
    */
   uint32 public constant WEEK = 120_960;
+
+  /**
+   * @dev Duration of a week in blocks on the CELO blockchain after the L2 transition (1 seconds per block)
+   */
+  uint32 public constant L2_WEEK = 604_800;
+
   /**
    * @dev Maximum allowable cliff period for token locks in weeks
    */
@@ -84,6 +92,31 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
    * @dev Total supply line of veMento
    */
   LibBrokenLine.BrokenLine public totalSupplyLine;
+
+  // *************** 
+  // New variables for L2 transition upgrade (4 slots)
+  // ***************
+  /**
+   * @dev L2 transtion block number
+   */
+  uint256 public l2Block;
+  /**
+   * @dev Address of the Mento Labs multisig
+   */
+  address public mentoLabsMultisig;
+  /**
+   * @dev L2 starting point week number used to move the first week after the L2 transition to the last week before the L2 transition
+   */
+  int256 public l2StartingPointWeek;
+  /**
+   * @dev Shift amount used after L2 transition to move the start of the epoch to 00-00 UTC Wednesday (approx)
+   */
+  uint32 public l2Shift;
+  /**
+   * @dev Flag to pause locking and governance
+   */
+  bool public paused;
+
   /**
    * @dev Emitted when create Lock with parameters (account, delegate, amount, slopePeriod, cliff)
    */
@@ -149,6 +182,11 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
     minSlopePeriod = _minSlopePeriod;
   }
 
+  modifier onlyMentoLabs() {
+    require(msg.sender == mentoLabsMultisig, "caller is not MentoLabs multisig");
+    _;
+  }
+
   /**
    * @notice Adds a new locking line for an account, initializing the lock with specified parameters.
    * @param account  Address for which tokens are being locked.
@@ -256,20 +294,29 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
 
   /**
    * @notice Calculates the week number for a given blocknumber
+   * @dev It takes L2 transition into account to calculate the week number consistently
    * @param ts block number
    * @return week number the block number belongs to
    */
   function roundTimestamp(uint32 ts) public view returns (uint32) {
+    require(!paused, "locking is paused");
+    
     if (ts < getEpochShift()) {
       return 0;
     }
-    uint32 shifted = ts - (getEpochShift());
-    return shifted / WEEK - uint32(startingPointWeek);
+
+    if (l2Block == 0 || ts < l2Block) {
+      uint32 shifted = ts - getEpochShift();
+      return shifted / WEEK - uint32(startingPointWeek);
+    } else {
+      uint32 shifted = ts - l2Shift;
+      return uint32(uint256(int256(uint256(shifted / L2_WEEK)) - l2StartingPointWeek));
+    }
   }
 
   /**
-   * @notice method returns the amount of blocks to shift locking epoch to.
-   * we move it to 00-00 UTC Wednesday (approx) by shifting 89964 blocks (CELO)
+   * @notice method returns the amount of blocks to shift locking epoch to on L1 CELO.
+   * we move it to 00-00 UTC Wednesday (approx) by shifting 89964 blocks 
    */
   function getEpochShift() internal view virtual returns (uint32) {
     return 89964;
@@ -351,6 +398,50 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
     uint32 time = roundTimestamp(blockNumber);
     updateTotalSupplyLine(time);
   }
+  /**
+    * @notice Sets the Mento Labs multisig address
+    * @param mentoLabsMultisig_ address of the Mento Labs multisig
+    */
+   */
+  function setMentoLabsMultisig(address mentoLabsMultisig_) external onlyOwner {
+    mentoLabsMultisig = mentoLabsMultisig_;
+  }
+  /**
+    * @notice Sets the L2 transition block number and pauses locking and governance
+    * @param blockNo block number of the L2 transition
+    */
+   */
+  function setL2TransitionBlock(uint256 blockNo) external onlyMentoLabs{
+
+    l2Block = blockNo;
+    paused = true;
+  }
+
+  /**
+    * @notice Sets the L2 shift amount
+    * @param _l2Shift shift amount that will be used after L2 transition
+   */
+  function setL2Shift(uint32 l2Shift_) external onlyMentoLabs{
+
+    l2Shift = l2Shift_;
+  }
+
+  /**
+    * @notice Sets the L2 starting point week number
+    * @param l2StartingPointWeek_ starting point week number that will be used after L2 transition
+   */
+  function setL2StartingPointWeek(int256 l2StartingPointWeek_) external onlyMentoLabs {
+
+    l2StartingPointWeek = l2StartingPointWeek_;
+  }
+
+  /**
+    * @notice Sets the paused flag
+    * @param paused_ flag to pause locking and governance
+   */
+  function setPaused(bool paused_) external onlyMentoLabs {
+    paused = paused_;
+  }
 
-  uint256[50] private __gap;
+  uint256[46] private __gap;
 }
diff --git a/test/fork/ForkTests.t.sol b/test/fork/ForkTests.t.sol
index d3ab0b48..ce1adbbc 100644
--- a/test/fork/ForkTests.t.sol
+++ b/test/fork/ForkTests.t.sol
@@ -44,6 +44,7 @@ import { BancorExchangeProviderForkTest } from "./BancorExchangeProviderForkTest
 import { GoodDollarTradingLimitsForkTest } from "./GoodDollar/TradingLimitsForkTest.sol";
 import { GoodDollarSwapForkTest } from "./GoodDollar/SwapForkTest.sol";
 import { GoodDollarExpansionForkTest } from "./GoodDollar/ExpansionForkTest.sol";
+import { LockingUpgradeForkTest } from "./upgrades/LockingUpgradeForkTest.sol";
 
 contract Alfajores_ChainForkTest is ChainForkTest(ALFAJORES_ID, 1, uints(15)) {}
 
@@ -116,3 +117,5 @@ contract Celo_GoodDollarTradingLimitsForkTest is GoodDollarTradingLimitsForkTest
 contract Celo_GoodDollarSwapForkTest is GoodDollarSwapForkTest(CELO_ID) {}
 
 contract Celo_GoodDollarExpansionForkTest is GoodDollarExpansionForkTest(CELO_ID) {}
+
+contract Celo_LockingUpgradeForkTest is LockingUpgradeForkTest(CELO_ID) {}
diff --git a/test/unit/governance/Locking/LockingTest.sol b/test/unit/governance/Locking/LockingTest.sol
index 45f7b6b7..f55f0356 100644
--- a/test/unit/governance/Locking/LockingTest.sol
+++ b/test/unit/governance/Locking/LockingTest.sol
@@ -31,4 +31,8 @@ contract LockingTest is GovernanceTest {
   function _incrementBlock(uint32 _amount) internal {
     locking.incrementBlock(_amount);
   }
+
+  function _reduceBlock(uint32 _amount) internal {
+    locking.reduceBlock(_amount);
+  }
 }
diff --git a/test/unit/governance/Locking/locking.upgrade.t.sol b/test/unit/governance/Locking/locking.upgrade.t.sol
new file mode 100644
index 00000000..f7a3ab47
--- /dev/null
+++ b/test/unit/governance/Locking/locking.upgrade.t.sol
@@ -0,0 +1,384 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+pragma solidity 0.8.18;
+// solhint-disable func-name-mixedcase, contract-name-camelcase
+
+import { LockingTest } from "./LockingTest.sol";
+
+import "forge-std/console.sol";
+
+contract Upgrade_LockingTest is LockingTest {
+
+  address public mentoLabs = makeAddr("MentoLabsMultisig");
+
+  uint32 public l1Day;
+  uint32 public l2Day;
+  uint32 public l1Week;
+  uint32 public l2Week;
+
+   function setUp() public override{
+    super.setUp();
+    l1Week = locking.WEEK();
+    l2Week = locking.L2_WEEK();
+    l1Day = l1Week / 7;
+    l2Day = l2Week / 7;
+  }
+
+  function test_initialSetup_shouldHaveCorrectValues() public view {
+    assertEq(locking.L2_WEEK(), 7 days);
+    assertEq(locking.mentoLabsMultisig(), address(0));
+    assertEq(locking.l2Block(), 0);
+    assertEq(locking.l2StartingPointWeek(), 0);
+    assertEq(locking.l2Shift(), 0);
+    assert(!locking.paused());
+  }
+
+  function test_setMentoLabsMultisig_whenCalledByNonOwner_shouldRevert() public {
+    vm.prank(alice);
+    vm.expectRevert("Ownable: caller is not the owner");
+    locking.setMentoLabsMultisig(mentoLabs);
+  }
+
+  function test_setMentoLabsMultisig_whenCalledByOwner_shouldSetMultisigAddress() public {
+    assertEq(locking.mentoLabsMultisig(), address(0));
+    
+    vm.prank(owner);
+    locking.setMentoLabsMultisig(mentoLabs);
+
+    assertEq(locking.mentoLabsMultisig(), mentoLabs);
+  }
+
+  modifier setMultisig() {
+    vm.prank(owner);
+    locking.setMentoLabsMultisig(mentoLabs);
+    _;
+  }
+
+  function test_setL2TransitionBlock_whenCalledByNonMentoMultisig_shouldRevert() public setMultisig {
+    vm.prank(alice);
+    vm.expectRevert("caller is not MentoLabs multisig");
+    locking.setL2TransitionBlock(block.number);
+  }
+
+  function test_setL2TransitionBlock_whenCalledByMentoMultisig_shouldSetL2BlockAndPause() public setMultisig {
+    assertEq(locking.l2Block(), 0);
+    assert(!locking.paused());
+    
+    uint32 blockNumber = uint32(block.number + 100);
+    
+    vm.prank(mentoLabs);
+    locking.setL2TransitionBlock(blockNumber);
+
+    assertEq(locking.l2Block(), blockNumber);
+    assert(locking.paused());
+  }
+
+  function test_setL2Shift_whenCalledByNonMentoMultisig_shouldRevert() public setMultisig {
+    vm.prank(alice);
+    vm.expectRevert("caller is not MentoLabs multisig");
+    locking.setL2Shift(100);
+  }
+
+  function test_setL2Shift_whenCalledByMentoMultisig_shouldSetL2BlockAndPause() public setMultisig {
+    assertEq(locking.l2Shift(), 0);
+    
+    vm.prank(mentoLabs);
+    locking.setL2Shift(100);
+
+    assertEq(locking.l2Shift(), 100);
+  }
+
+  function test_setL2StartingPointWeek_whenCalledByNonMentoMultisig_shouldRevert() public setMultisig {
+    vm.prank(alice);
+    vm.expectRevert("caller is not MentoLabs multisig");
+    locking.setL2StartingPointWeek(100);
+  }
+
+  function test_setL2StartingPointWeek_whenCalledByMentoMultisig_shouldSetL2BlockAndPause() public setMultisig {
+    assertEq(locking.l2StartingPointWeek(), 0);
+    
+    
+    vm.prank(mentoLabs);
+    locking.setL2StartingPointWeek(100);
+
+    assertEq(locking.l2StartingPointWeek(), 100);
+  }
+  function test_setPaused_whenCalledByNonMentoMultisig_shouldRevert() public setMultisig {
+    vm.prank(alice);
+    vm.expectRevert("caller is not MentoLabs multisig");
+    locking.setPaused(true);
+  }
+
+  function test_setPaused_whenCalledByMentoMultisig_shouldPauseContracts() public setMultisig {
+    assert(!locking.paused());
+
+    mentoToken.mint(alice, 1000000e18);
+    
+    vm.prank(mentoLabs);
+    locking.setPaused(true);
+
+
+    assert(locking.paused());
+
+    vm.expectRevert("locking is paused");
+    vm.prank(alice);
+    locking.lock(alice, bob, 1000e18, 5, 5);
+
+    vm.expectRevert("locking is paused");
+    vm.prank(alice);
+    locking.withdraw();
+
+
+    vm.prank(mentoLabs);
+    locking.setPaused(false);
+
+    assert(!locking.paused());
+  }
+
+  modifier l2LockingSetup(uint32 advanceWeeks_, uint32 startingPointWeek_, uint32 l1Shift_) {
+    vm.prank(owner);
+    locking.setMentoLabsMultisig(mentoLabs);
+
+    _incrementBlock(l1Week * advanceWeeks_);
+
+    locking.setStatingPointWeek(startingPointWeek_);
+    locking.setEpochShift(l1Shift_);
+
+    vm.prank(mentoLabs);
+    locking.setL2TransitionBlock(block.number);
+
+    vm.prank(mentoLabs);
+    locking.setPaused(false);
+
+    _;
+  }
+
+  function test_getWeek_whenShiftAndStartingPointIs0_shouldReturnCorrectWeekNo() public l2LockingSetup(8, 0, 0) {
+
+    // 2 + 8 weeks = 10 weeks on l1 = 2 weeks on l2
+    assertEq(locking.getWeek(), 2);
+    assertEq(locking.l2BlockTillNextPeriod(), l2Week);
+
+    _incrementBlock(l2Day * 3);
+
+    assertEq(locking.getWeek(), 2);
+    assertEq(locking.l2BlockTillNextPeriod(), l2Day * 4);
+
+    _incrementBlock(l2Day * 5);
+
+    assertEq(locking.getWeek(), 3);
+    assertEq(locking.l2BlockTillNextPeriod(), l2Day * 6);
+
+    _incrementBlock(l2Day * 8);
+
+    assertEq(locking.getWeek(), 4);
+    assertEq(locking.l2BlockTillNextPeriod(), l2Day * 5);
+  }
+
+  function test_getWeek_whenL2StartingPointIsPositive_shouldReturnCorrectWeekNo() public l2LockingSetup(198, 190, 0) {
+    // l1 week no = 198 + 2 - (190) = 10
+    uint32 l1WeekNo = 10;
+    // l2 week no = (198 + 2) / 5 = 40
+    assertEq(locking.getWeek(), 40);
+    assertEq(locking.l2BlockTillNextPeriod(), l2Week);
+
+    // l2WeekNo - l1WeekNo = 40 - 10 = 30
+    vm.prank(mentoLabs);
+    locking.setL2StartingPointWeek(30);
+
+    // after the L2 starting point week is set, the week should be equal to the l1 week no
+    assertEq(locking.getWeek(), l1WeekNo);
+  }
+
+  function test_getWeek_whenL2StartingPointIsNegative_shouldReturnCorrectWeekNo() public l2LockingSetup(18,0,0) {
+    // l1 week no = 18 + 2 - 0 = 20
+    uint32 l1WeekNo = 20;
+    // l2 week no = (18 + 2) / 5 = 4
+    assertEq(locking.getWeek(), 4);
+    assertEq(locking.l2BlockTillNextPeriod(), l2Week);
+
+    // l2WeekNo - l1WeekNo = 4 - 20 = -16
+    vm.prank(mentoLabs);
+    locking.setL2StartingPointWeek(-16);
+
+    // after the L2 starting point week is set, the week should be equal to the l1 week no
+    assertEq(locking.getWeek(), l1WeekNo);
+  }
+
+  function test_getWeek_whenShiftIsPositive_shouldReturnCorrectWeekNo() public l2LockingSetup(18, 5, l1Day * 3) {
+    // l1 week no = 18 + 2 - 5 - 1 = 19
+    uint32 l1WeekNo = 14;
+
+    // l2 week no = (18 + 2) / 5 = 4
+    assertEq(locking.getWeek(), 4);
+    assertEq(locking.l2BlockTillNextPeriod(), l2Week);
+
+    // l2WeekNo - l1WeekNo = 4 - 14 - 1 = -11
+    vm.prank(mentoLabs);
+    locking.setL2StartingPointWeek(-11);
+
+    vm.prank(mentoLabs);
+    locking.setL2Shift(l2Day * 3);
+
+    // after the L2 starting point week and l2Shift are set, the timing should be equal to the l1 timing
+    assertEq(locking.getWeek(), l1WeekNo);
+    assertEq(locking.l2BlockTillNextPeriod(), l2Day * 3);
+
+    _incrementBlock(l2Day);
+
+    assertEq(locking.getWeek(), l1WeekNo);
+    assertEq(locking.l2BlockTillNextPeriod(), l2Day * 2);
+
+    _incrementBlock(l2Day);
+
+    assertEq(locking.getWeek(), l1WeekNo);
+    assertEq(locking.l2BlockTillNextPeriod(), l2Day);
+
+
+    _incrementBlock(l2Day);
+
+    assertEq(locking.getWeek(), l1WeekNo + 1);
+    assertEq(locking.l2BlockTillNextPeriod(), l2Week);
+  }
+
+  function test_totalSupply_whenCalledAfterL2Transition_shouldReturnCorrectValues() public setMultisig {
+    mentoToken.mint(alice, 1000000e18);
+
+    // week no: 20
+    _incrementBlock(l1Week * 18);
+
+    vm.prank(alice);
+    locking.lock(alice, alice, 1000e18, 104, 0);
+
+    // week no: 40
+    _incrementBlock(l1Week * 20);
+
+
+    uint256 totalSupplyL1W40 = locking.totalSupply();
+    uint256 pastTotalSupplyL1W30= locking.getPastTotalSupply(locking.blockNumberMocked() - l1Week * 10);
+
+    // week no: 60
+    _incrementBlock(l1Week * 20);
+
+    uint256 totalSupplyL1W60 = locking.totalSupply();
+    uint256 pastTotalSupplyL1W50= locking.getPastTotalSupply(locking.blockNumberMocked() - l1Week * 10);
+
+    // roll back to week 40
+    _reduceBlock(l1Week * 20);
+    
+
+    vm.prank(mentoLabs);
+    locking.setL2TransitionBlock(l1Week * 40);
+
+
+    vm.prank(mentoLabs);
+    locking.setPaused(false);
+
+    // 8 - 40 = -32
+    vm.prank(mentoLabs);
+    locking.setL2StartingPointWeek(-32);
+
+    assertEq(locking.totalSupply(), totalSupplyL1W40);
+    assertEq(locking.getPastTotalSupply(locking.blockNumberMocked() - l1Week * 10), pastTotalSupplyL1W30);
+
+    // week no: 60
+    _incrementBlock(l2Week * 20);
+    assertEq(locking.totalSupply(), totalSupplyL1W60);
+    assertEq(locking.getPastTotalSupply(locking.blockNumberMocked() - l2Week * 10), pastTotalSupplyL1W50);
+  }
+
+  function test_balanceOfAndGetVotes_whenCalledAfterL2Transition_shouldReturnCorrectValues() public setMultisig {
+    mentoToken.mint(alice, 1000000e18);
+
+    //  week no: 20
+    _incrementBlock(l1Week * 18);
+
+    vm.prank(alice);
+    locking.lock(alice, alice, 1000e18, 104, 0);
+
+    // week no: 40
+    _incrementBlock(l1Week * 20);
+
+
+    uint256 balanceOfL1W40 = locking.balanceOf(alice);
+    uint256 votesL1W40 = locking.getVotes(alice);
+    uint256 pastVotesL1W30 = locking.getPastVotes(alice, locking.blockNumberMocked() - l1Week * 10);
+
+    // week no: 60
+    _incrementBlock(l1Week * 20);
+
+    uint256 balanceOfL1W60 = locking.balanceOf(alice);
+    uint256 votesL1W60 = locking.getVotes(alice);
+    uint256 pastVotesL1W50 = locking.getPastVotes(alice, locking.blockNumberMocked() - l1Week * 10);
+
+    // roll back to week 40
+    _reduceBlock(l1Week * 20);
+    
+
+    vm.prank(mentoLabs);
+    locking.setL2TransitionBlock(l1Week * 40);
+
+
+    vm.prank(mentoLabs);
+    locking.setPaused(false);
+
+    // 8 - 40 = -32
+    vm.prank(mentoLabs);
+    locking.setL2StartingPointWeek(-32);
+
+    assertEq(locking.balanceOf(alice), balanceOfL1W40);
+    assertEq(locking.getVotes(alice), votesL1W40);
+    assertEq(locking.getPastVotes(alice, locking.blockNumberMocked() - l1Week * 10), pastVotesL1W30);
+
+    // week no: 60
+    _incrementBlock(l2Week * 20);
+    assertEq(locking.balanceOf(alice), balanceOfL1W60);
+    assertEq(locking.getVotes(alice), votesL1W60);
+    assertEq(locking.getPastVotes(alice, locking.blockNumberMocked() - l2Week * 10), pastVotesL1W50);
+  }
+
+  function test_lockedAndWithdrawable_whenCalledAfterL2Transition_shouldReturnCorrectValues() public setMultisig {
+    mentoToken.mint(alice, 1000000e18);
+
+    //  week no: 20
+    _incrementBlock(l1Week * 18);
+
+    vm.prank(alice);
+    locking.lock(alice, alice, 1000e18, 104, 0);
+
+    // week no: 40
+    _incrementBlock(l1Week * 20);
+
+
+    uint256 lockedL1W40 = locking.locked(alice);
+    uint256 withdrawableL1W40 = locking.getAvailableForWithdraw(alice);
+
+    // week no: 60
+    _incrementBlock(l1Week * 20);
+
+    uint256 lockedL1W60 = locking.locked(alice);
+    uint256 withdrawableL1W60 = locking.getAvailableForWithdraw(alice); 
+
+    // roll back to week 40
+    _reduceBlock(l1Week * 20);
+    
+
+    vm.prank(mentoLabs);
+    locking.setL2TransitionBlock(l1Week * 40);
+
+    vm.prank(mentoLabs);
+    locking.setPaused(false);
+
+    // 8 - 40 = -32
+    vm.prank(mentoLabs);
+    locking.setL2StartingPointWeek(-32);
+
+    assertEq(locking.locked(alice), lockedL1W40);
+    assertEq(locking.getAvailableForWithdraw(alice), withdrawableL1W40);
+
+    // week no: 60
+    _incrementBlock(l2Week * 20);
+    assertEq(locking.locked(alice), lockedL1W60);
+    assertEq(locking.getAvailableForWithdraw(alice), withdrawableL1W60);
+  }
+ 
+}
diff --git a/test/utils/harnesses/LockingHarness.sol b/test/utils/harnesses/LockingHarness.sol
index 440708f7..e4e3521c 100644
--- a/test/utils/harnesses/LockingHarness.sol
+++ b/test/utils/harnesses/LockingHarness.sol
@@ -11,6 +11,10 @@ contract LockingHarness is Locking {
     blockNumberMocked = blockNumberMocked + _amount;
   }
 
+  function reduceBlock(uint32 _amount) external {
+    blockNumberMocked = blockNumberMocked - _amount;
+  }
+
   function getBlockNumber() internal view override returns (uint32) {
     return blockNumberMocked;
   }
@@ -45,6 +49,16 @@ contract LockingHarness is Locking {
 
   function blockTillNextPeriod() external view returns (uint256) {
     uint256 currentWeek = this.getWeek();
-    return (WEEK * (currentWeek + 1)) + getEpochShift() - getBlockNumber();
+    return (WEEK * (currentWeek + startingPointWeek + 1)) + getEpochShift() - getBlockNumber();
   }
+
+  function l2BlockTillNextPeriod() external view returns (uint256) {
+    uint256 currentWeek = this.getWeek();
+    return (L2_WEEK * uint256(int(currentWeek) + l2StartingPointWeek + 1)) + l2Shift - getBlockNumber();
+  }
+
+
+  function setStatingPointWeek(uint32 _week) external {
+    startingPointWeek = _week;
+  } 
 }

From 08431e3b4cf9cd2d3760aad88c1dfeb35ad89e02 Mon Sep 17 00:00:00 2001
From: baroooo <baranseltekin@gmail.com>
Date: Mon, 18 Nov 2024 14:13:38 +0100
Subject: [PATCH 02/16] feat: fork integration test for l2 upgrade

---
 contracts/governance/locking/LockingBase.sol  |   6 +-
 test/fork/upgrades/LockingUpgradeForkTest.sol | 178 ++++++++++++++++++
 .../governance/Locking/locking.upgrade.t.sol  |   4 +-
 3 files changed, 183 insertions(+), 5 deletions(-)
 create mode 100644 test/fork/upgrades/LockingUpgradeForkTest.sol

diff --git a/contracts/governance/locking/LockingBase.sol b/contracts/governance/locking/LockingBase.sol
index 24211440..0c6b3bdc 100644
--- a/contracts/governance/locking/LockingBase.sol
+++ b/contracts/governance/locking/LockingBase.sol
@@ -401,16 +401,16 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
   /**
     * @notice Sets the Mento Labs multisig address
     * @param mentoLabsMultisig_ address of the Mento Labs multisig
+    *
     */
-   */
   function setMentoLabsMultisig(address mentoLabsMultisig_) external onlyOwner {
     mentoLabsMultisig = mentoLabsMultisig_;
   }
   /**
     * @notice Sets the L2 transition block number and pauses locking and governance
     * @param blockNo block number of the L2 transition
+    *
     */
-   */
   function setL2TransitionBlock(uint256 blockNo) external onlyMentoLabs{
 
     l2Block = blockNo;
@@ -419,7 +419,7 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
 
   /**
     * @notice Sets the L2 shift amount
-    * @param _l2Shift shift amount that will be used after L2 transition
+    * @param l2Shift_ shift amount that will be used after L2 transition
    */
   function setL2Shift(uint32 l2Shift_) external onlyMentoLabs{
 
diff --git a/test/fork/upgrades/LockingUpgradeForkTest.sol b/test/fork/upgrades/LockingUpgradeForkTest.sol
new file mode 100644
index 00000000..c1fb6d3e
--- /dev/null
+++ b/test/fork/upgrades/LockingUpgradeForkTest.sol
@@ -0,0 +1,178 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+pragma solidity ^0.8;
+
+import { BaseForkTest } from "../BaseForkTest.sol";
+import { Locking } from "contracts/governance/locking/Locking.sol";
+import { GovernanceFactory } from "contracts/governance/GovernanceFactory.sol";
+import { ProxyAdmin } from "openzeppelin-contracts-next/contracts/proxy/transparent/ProxyAdmin.sol";
+import { ITransparentUpgradeableProxy } from "openzeppelin-contracts-next/contracts/proxy/transparent/TransparentUpgradeableProxy.sol";
+import "forge-std/console.sol";
+
+// used to avoid stack too deep error
+struct Balances {
+  uint256 totalSupply;
+  uint256 pastTotalSupply;
+  uint256 balance1;
+  uint256 balance2;
+  uint256 votingPower1;
+  uint256 votingPower2;
+  uint256 pastVotingPower1;
+  uint256 pastVotingPower2;
+  uint256 lockedBalance1;
+  uint256 lockedBalance2;
+  uint256 withdrawable1;
+  uint256 withdrawable2;
+}
+contract LockingUpgradeForkTest is BaseForkTest {
+  // airdrop claimers from the mainnet
+  address public constant AIRDROP_CLAIMER_1 = 0x3152eE4a18ee3209524F9071A6BcAdA098f19838;
+  address public constant AIRDROP_CLAIMER_2 = 0x44EB9Bf2D6B161499f1b706c331aa2Ba1d5069c7;
+
+  uint256 public constant L1_WEEK = 7 days / 5;
+  uint256 public constant L2_WEEK = 7 days;
+
+  GovernanceFactory public governanceFactory = GovernanceFactory(0xee6CE2dbe788dFC38b8F583Da86cB9caf2C8cF5A);
+  ProxyAdmin public proxyAdmin;
+  Locking public locking;
+  address public timelockController;
+  address public newLockingImplementation;
+
+  address public mentoLabsMultisig = makeAddr("mentoLabsMultisig");
+
+  constructor(uint256 _chainId) BaseForkTest(_chainId) {}
+
+  function setUp() public virtual override {
+    super.setUp();
+    proxyAdmin = governanceFactory.proxyAdmin();
+    locking = governanceFactory.locking();
+    timelockController = address(governanceFactory.governanceTimelock());
+
+    newLockingImplementation = address(new Locking());
+    vm.prank(timelockController);
+    proxyAdmin.upgrade(ITransparentUpgradeableProxy(address(locking)), newLockingImplementation);
+  }
+
+  function test_upgrade() public {
+    Balances memory beforeBalances;
+    Balances memory afterBalances;
+
+    vm.prank(timelockController);
+    locking.setMentoLabsMultisig(mentoLabsMultisig);
+
+    // THU Nov-07-2024 12:00:23 AM +UTC
+    vm.roll(28653031);
+    vm.warp(1730937623);
+
+    // move 3 weeks forward on L1
+    moveDays(3 * 7, true, false);
+
+    uint256 weekNoBefore = locking.getWeek();
+
+    beforeBalances.totalSupply = locking.totalSupply();
+    beforeBalances.pastTotalSupply = locking.getPastTotalSupply(block.number - 3 * L1_WEEK);
+    beforeBalances.balance1 = locking.balanceOf(AIRDROP_CLAIMER_1);
+    beforeBalances.balance2 = locking.balanceOf(AIRDROP_CLAIMER_2);
+    beforeBalances.votingPower1 = locking.getVotes(AIRDROP_CLAIMER_1);
+    beforeBalances.votingPower2 = locking.getVotes(AIRDROP_CLAIMER_2);
+    beforeBalances.pastVotingPower1 = locking.getPastVotes(AIRDROP_CLAIMER_1, block.number - 3 * L1_WEEK);
+    beforeBalances.pastVotingPower2 = locking.getPastVotes(AIRDROP_CLAIMER_2, block.number - 3 * L1_WEEK);
+    beforeBalances.lockedBalance1 = locking.locked(AIRDROP_CLAIMER_1);
+    beforeBalances.lockedBalance2 = locking.locked(AIRDROP_CLAIMER_2);
+    beforeBalances.withdrawable1 = locking.getAvailableForWithdraw(AIRDROP_CLAIMER_1);
+    beforeBalances.withdrawable2 = locking.getAvailableForWithdraw(AIRDROP_CLAIMER_2);
+
+    // move 5 weeks forward on L1
+    moveDays(5 * 7, true, false);
+
+    afterBalances.totalSupply = locking.totalSupply();
+    afterBalances.pastTotalSupply = locking.getPastTotalSupply(block.number - 3 * L1_WEEK);
+    afterBalances.balance1 = locking.balanceOf(AIRDROP_CLAIMER_1);
+    afterBalances.balance2 = locking.balanceOf(AIRDROP_CLAIMER_2);
+    afterBalances.votingPower1 = locking.getVotes(AIRDROP_CLAIMER_1);
+    afterBalances.votingPower2 = locking.getVotes(AIRDROP_CLAIMER_2);
+    afterBalances.pastVotingPower1 = locking.getPastVotes(AIRDROP_CLAIMER_1, block.number - 3 * L1_WEEK);
+    afterBalances.pastVotingPower2 = locking.getPastVotes(AIRDROP_CLAIMER_2, block.number - 3 * L1_WEEK);
+    afterBalances.lockedBalance1 = locking.locked(AIRDROP_CLAIMER_1);
+    afterBalances.lockedBalance2 = locking.locked(AIRDROP_CLAIMER_2);
+    afterBalances.withdrawable1 = locking.getAvailableForWithdraw(AIRDROP_CLAIMER_1);
+    afterBalances.withdrawable2 = locking.getAvailableForWithdraw(AIRDROP_CLAIMER_2);
+
+    // move 35 days backward on L1
+    moveDays(35, false, false);
+
+    uint256 blocksTillNextWeek = ((L1_WEEK * (locking.getWeek() + locking.startingPointWeek() + 1)) + 89964) -
+      block.number;
+
+    // simulate L2 upgrade
+    vm.prank(mentoLabsMultisig);
+    locking.setL2TransitionBlock(block.number);
+
+    vm.prank(mentoLabsMultisig);
+    locking.setL2StartingPointWeek(20);
+
+    vm.prank(mentoLabsMultisig);
+    locking.setL2Shift(507776);
+
+    vm.prank(mentoLabsMultisig);
+    locking.setPaused(false);
+
+    uint256 blocksTillNextWeek2 = ((L2_WEEK * (locking.getWeek() + uint256(locking.l2StartingPointWeek()) + 1)) +
+      507776) - block.number;
+
+    // if the shift number is correct, the number of blocks till the next week should be 5 times the previous number
+    assertEq(blocksTillNextWeek2, 5 * blocksTillNextWeek);
+
+    assertEq(locking.getWeek(), weekNoBefore);
+    assertEq(locking.totalSupply(), beforeBalances.totalSupply);
+    // the past values should be calculated using the L1 week value
+    assertEq(locking.getPastTotalSupply(block.number - 3 * L1_WEEK), beforeBalances.pastTotalSupply);
+    assertEq(locking.balanceOf(AIRDROP_CLAIMER_1), beforeBalances.balance1);
+    assertEq(locking.balanceOf(AIRDROP_CLAIMER_2), beforeBalances.balance2);
+    assertEq(locking.getVotes(AIRDROP_CLAIMER_1), beforeBalances.votingPower1);
+    assertEq(locking.getVotes(AIRDROP_CLAIMER_2), beforeBalances.votingPower2);
+    assertEq(locking.getPastVotes(AIRDROP_CLAIMER_1, block.number - 3 * L1_WEEK), beforeBalances.pastVotingPower1);
+    assertEq(locking.getPastVotes(AIRDROP_CLAIMER_2, block.number - 3 * L1_WEEK), beforeBalances.pastVotingPower2);
+    assertEq(locking.locked(AIRDROP_CLAIMER_1), beforeBalances.lockedBalance1);
+    assertEq(locking.locked(AIRDROP_CLAIMER_2), beforeBalances.lockedBalance2);
+    assertEq(locking.getAvailableForWithdraw(AIRDROP_CLAIMER_1), beforeBalances.withdrawable1);
+    assertEq(locking.getAvailableForWithdraw(AIRDROP_CLAIMER_2), beforeBalances.withdrawable2);
+
+    // // move 5 weeks forward on L2
+    moveDays(5 * 7, true, true);
+
+    assertEq(locking.getWeek(), weekNoBefore + 5);
+
+    blocksTillNextWeek2 =
+      ((L2_WEEK * (locking.getWeek() + uint256(locking.l2StartingPointWeek()) + 1)) + 507776) -
+      block.number;
+
+    assertEq(blocksTillNextWeek2, 5 * blocksTillNextWeek);
+    assertEq(locking.totalSupply(), afterBalances.totalSupply);
+    assertEq(locking.getPastTotalSupply(block.number - 3 * L2_WEEK), afterBalances.pastTotalSupply);
+    assertEq(locking.balanceOf(AIRDROP_CLAIMER_1), afterBalances.balance1);
+    assertEq(locking.balanceOf(AIRDROP_CLAIMER_2), afterBalances.balance2);
+    assertEq(locking.getVotes(AIRDROP_CLAIMER_1), afterBalances.votingPower1);
+    assertEq(locking.getVotes(AIRDROP_CLAIMER_2), afterBalances.votingPower2);
+    assertEq(locking.getPastVotes(AIRDROP_CLAIMER_1, block.number - 3 * L2_WEEK), afterBalances.pastVotingPower1);
+    assertEq(locking.getPastVotes(AIRDROP_CLAIMER_2, block.number - 3 * L2_WEEK), afterBalances.pastVotingPower2);
+    assertEq(locking.locked(AIRDROP_CLAIMER_1), afterBalances.lockedBalance1);
+    assertEq(locking.locked(AIRDROP_CLAIMER_2), afterBalances.lockedBalance2);
+    assertEq(locking.getAvailableForWithdraw(AIRDROP_CLAIMER_1), afterBalances.withdrawable1);
+    assertEq(locking.getAvailableForWithdraw(AIRDROP_CLAIMER_2), afterBalances.withdrawable2);
+  }
+
+  // move days forward or backward on L1 or L2
+  function moveDays(uint256 day, bool forward, bool isL2) public {
+    uint256 ts = vm.getBlockTimestamp();
+    uint256 height = vm.getBlockNumber();
+
+    uint256 newTs = forward ? ts + day * 1 days : ts - day * 1 days;
+    uint256 blockChange = isL2 ? (day * 1 days) : ((day * 1 days) / 5);
+    uint256 newHeight = forward ? height + blockChange : height - blockChange;
+    vm.warp(newTs);
+    vm.roll(newHeight);
+
+    ts = vm.getBlockTimestamp();
+    height = vm.getBlockNumber();
+  }
+}
diff --git a/test/unit/governance/Locking/locking.upgrade.t.sol b/test/unit/governance/Locking/locking.upgrade.t.sol
index f7a3ab47..9c71b423 100644
--- a/test/unit/governance/Locking/locking.upgrade.t.sol
+++ b/test/unit/governance/Locking/locking.upgrade.t.sol
@@ -17,8 +17,8 @@ contract Upgrade_LockingTest is LockingTest {
 
    function setUp() public override{
     super.setUp();
-    l1Week = locking.WEEK();
-    l2Week = locking.L2_WEEK();
+    l1Week = 7 days / 5;
+    l2Week = 7 days;
     l1Day = l1Week / 7;
     l2Day = l2Week / 7;
   }

From 34cb12987ae07640a97dc61824ab0b624146b152 Mon Sep 17 00:00:00 2001
From: baroooo <baranseltekin@gmail.com>
Date: Mon, 18 Nov 2024 15:29:05 +0100
Subject: [PATCH 03/16] chore: refactor fork test

---
 test/fork/upgrades/LockingUpgradeForkTest.sol | 192 ++++++++++--------
 1 file changed, 102 insertions(+), 90 deletions(-)

diff --git a/test/fork/upgrades/LockingUpgradeForkTest.sol b/test/fork/upgrades/LockingUpgradeForkTest.sol
index c1fb6d3e..db6e8dfa 100644
--- a/test/fork/upgrades/LockingUpgradeForkTest.sol
+++ b/test/fork/upgrades/LockingUpgradeForkTest.sol
@@ -6,10 +6,9 @@ import { Locking } from "contracts/governance/locking/Locking.sol";
 import { GovernanceFactory } from "contracts/governance/GovernanceFactory.sol";
 import { ProxyAdmin } from "openzeppelin-contracts-next/contracts/proxy/transparent/ProxyAdmin.sol";
 import { ITransparentUpgradeableProxy } from "openzeppelin-contracts-next/contracts/proxy/transparent/TransparentUpgradeableProxy.sol";
-import "forge-std/console.sol";
 
 // used to avoid stack too deep error
-struct Balances {
+struct LockingSnapshot {
   uint256 totalSupply;
   uint256 pastTotalSupply;
   uint256 balance1;
@@ -23,6 +22,7 @@ struct Balances {
   uint256 withdrawable1;
   uint256 withdrawable2;
 }
+
 contract LockingUpgradeForkTest is BaseForkTest {
   // airdrop claimers from the mainnet
   address public constant AIRDROP_CLAIMER_1 = 0x3152eE4a18ee3209524F9071A6BcAdA098f19838;
@@ -50,129 +50,141 @@ contract LockingUpgradeForkTest is BaseForkTest {
     newLockingImplementation = address(new Locking());
     vm.prank(timelockController);
     proxyAdmin.upgrade(ITransparentUpgradeableProxy(address(locking)), newLockingImplementation);
-  }
-
-  function test_upgrade() public {
-    Balances memory beforeBalances;
-    Balances memory afterBalances;
 
     vm.prank(timelockController);
     locking.setMentoLabsMultisig(mentoLabsMultisig);
 
-    // THU Nov-07-2024 12:00:23 AM +UTC
+    // THU Nov-07-2024 00:00:23 +UTC
     vm.roll(28653031);
     vm.warp(1730937623);
+  }
+
+  function test_blockNoDependentCalculations_afterL2Transition_shouldWorkAsBefore() public {
+    LockingSnapshot memory beforeSnapshot;
+    LockingSnapshot memory afterSnapshot;
 
     // move 3 weeks forward on L1
-    moveDays(3 * 7, true, false);
+    _moveDays({ day: 3 * 7, forward: true, isL2: false });
 
     uint256 weekNoBefore = locking.getWeek();
 
-    beforeBalances.totalSupply = locking.totalSupply();
-    beforeBalances.pastTotalSupply = locking.getPastTotalSupply(block.number - 3 * L1_WEEK);
-    beforeBalances.balance1 = locking.balanceOf(AIRDROP_CLAIMER_1);
-    beforeBalances.balance2 = locking.balanceOf(AIRDROP_CLAIMER_2);
-    beforeBalances.votingPower1 = locking.getVotes(AIRDROP_CLAIMER_1);
-    beforeBalances.votingPower2 = locking.getVotes(AIRDROP_CLAIMER_2);
-    beforeBalances.pastVotingPower1 = locking.getPastVotes(AIRDROP_CLAIMER_1, block.number - 3 * L1_WEEK);
-    beforeBalances.pastVotingPower2 = locking.getPastVotes(AIRDROP_CLAIMER_2, block.number - 3 * L1_WEEK);
-    beforeBalances.lockedBalance1 = locking.locked(AIRDROP_CLAIMER_1);
-    beforeBalances.lockedBalance2 = locking.locked(AIRDROP_CLAIMER_2);
-    beforeBalances.withdrawable1 = locking.getAvailableForWithdraw(AIRDROP_CLAIMER_1);
-    beforeBalances.withdrawable2 = locking.getAvailableForWithdraw(AIRDROP_CLAIMER_2);
+    // Take snapshot 3 weeks after Nov 07
+    beforeSnapshot = _takeSnapshot(AIRDROP_CLAIMER_1, AIRDROP_CLAIMER_2);
 
     // move 5 weeks forward on L1
-    moveDays(5 * 7, true, false);
-
-    afterBalances.totalSupply = locking.totalSupply();
-    afterBalances.pastTotalSupply = locking.getPastTotalSupply(block.number - 3 * L1_WEEK);
-    afterBalances.balance1 = locking.balanceOf(AIRDROP_CLAIMER_1);
-    afterBalances.balance2 = locking.balanceOf(AIRDROP_CLAIMER_2);
-    afterBalances.votingPower1 = locking.getVotes(AIRDROP_CLAIMER_1);
-    afterBalances.votingPower2 = locking.getVotes(AIRDROP_CLAIMER_2);
-    afterBalances.pastVotingPower1 = locking.getPastVotes(AIRDROP_CLAIMER_1, block.number - 3 * L1_WEEK);
-    afterBalances.pastVotingPower2 = locking.getPastVotes(AIRDROP_CLAIMER_2, block.number - 3 * L1_WEEK);
-    afterBalances.lockedBalance1 = locking.locked(AIRDROP_CLAIMER_1);
-    afterBalances.lockedBalance2 = locking.locked(AIRDROP_CLAIMER_2);
-    afterBalances.withdrawable1 = locking.getAvailableForWithdraw(AIRDROP_CLAIMER_1);
-    afterBalances.withdrawable2 = locking.getAvailableForWithdraw(AIRDROP_CLAIMER_2);
-
-    // move 35 days backward on L1
-    moveDays(35, false, false);
-
-    uint256 blocksTillNextWeek = ((L1_WEEK * (locking.getWeek() + locking.startingPointWeek() + 1)) + 89964) -
-      block.number;
-
-    // simulate L2 upgrade
-    vm.prank(mentoLabsMultisig);
-    locking.setL2TransitionBlock(block.number);
+    _moveDays({ day: 5 * 7, forward: true, isL2: false });
 
-    vm.prank(mentoLabsMultisig);
-    locking.setL2StartingPointWeek(20);
+    // Take snapshot 8 weeks after Nov 07
+    afterSnapshot = _takeSnapshot(AIRDROP_CLAIMER_1, AIRDROP_CLAIMER_2);
 
-    vm.prank(mentoLabsMultisig);
-    locking.setL2Shift(507776);
+    // move 5 weeks backward on L1
+    _moveDays({ day: 5 * 7, forward: false, isL2: false });
 
-    vm.prank(mentoLabsMultisig);
-    locking.setPaused(false);
+    uint256 blocksTillNextWeekL1 = _calculateBlocksTillNextWeek({ isL2: false });
+
+    _simulateL2Upgrade();
 
-    uint256 blocksTillNextWeek2 = ((L2_WEEK * (locking.getWeek() + uint256(locking.l2StartingPointWeek()) + 1)) +
-      507776) - block.number;
+    uint256 blocksTillNextWeekL2 = _calculateBlocksTillNextWeek({ isL2: true });
 
     // if the shift number is correct, the number of blocks till the next week should be 5 times the previous number
-    assertEq(blocksTillNextWeek2, 5 * blocksTillNextWeek);
+    assertEq(blocksTillNextWeekL2, 5 * blocksTillNextWeekL1);
 
     assertEq(locking.getWeek(), weekNoBefore);
-    assertEq(locking.totalSupply(), beforeBalances.totalSupply);
+    assertEq(locking.totalSupply(), beforeSnapshot.totalSupply);
     // the past values should be calculated using the L1 week value
-    assertEq(locking.getPastTotalSupply(block.number - 3 * L1_WEEK), beforeBalances.pastTotalSupply);
-    assertEq(locking.balanceOf(AIRDROP_CLAIMER_1), beforeBalances.balance1);
-    assertEq(locking.balanceOf(AIRDROP_CLAIMER_2), beforeBalances.balance2);
-    assertEq(locking.getVotes(AIRDROP_CLAIMER_1), beforeBalances.votingPower1);
-    assertEq(locking.getVotes(AIRDROP_CLAIMER_2), beforeBalances.votingPower2);
-    assertEq(locking.getPastVotes(AIRDROP_CLAIMER_1, block.number - 3 * L1_WEEK), beforeBalances.pastVotingPower1);
-    assertEq(locking.getPastVotes(AIRDROP_CLAIMER_2, block.number - 3 * L1_WEEK), beforeBalances.pastVotingPower2);
-    assertEq(locking.locked(AIRDROP_CLAIMER_1), beforeBalances.lockedBalance1);
-    assertEq(locking.locked(AIRDROP_CLAIMER_2), beforeBalances.lockedBalance2);
-    assertEq(locking.getAvailableForWithdraw(AIRDROP_CLAIMER_1), beforeBalances.withdrawable1);
-    assertEq(locking.getAvailableForWithdraw(AIRDROP_CLAIMER_2), beforeBalances.withdrawable2);
-
-    // // move 5 weeks forward on L2
-    moveDays(5 * 7, true, true);
+    assertEq(locking.getPastTotalSupply(block.number - 3 * L1_WEEK), beforeSnapshot.pastTotalSupply);
+    assertEq(locking.balanceOf(AIRDROP_CLAIMER_1), beforeSnapshot.balance1);
+    assertEq(locking.balanceOf(AIRDROP_CLAIMER_2), beforeSnapshot.balance2);
+    assertEq(locking.getVotes(AIRDROP_CLAIMER_1), beforeSnapshot.votingPower1);
+    assertEq(locking.getVotes(AIRDROP_CLAIMER_2), beforeSnapshot.votingPower2);
+    assertEq(locking.getPastVotes(AIRDROP_CLAIMER_1, block.number - 3 * L1_WEEK), beforeSnapshot.pastVotingPower1);
+    assertEq(locking.getPastVotes(AIRDROP_CLAIMER_2, block.number - 3 * L1_WEEK), beforeSnapshot.pastVotingPower2);
+    assertEq(locking.locked(AIRDROP_CLAIMER_1), beforeSnapshot.lockedBalance1);
+    assertEq(locking.locked(AIRDROP_CLAIMER_2), beforeSnapshot.lockedBalance2);
+    assertEq(locking.getAvailableForWithdraw(AIRDROP_CLAIMER_1), beforeSnapshot.withdrawable1);
+    assertEq(locking.getAvailableForWithdraw(AIRDROP_CLAIMER_2), beforeSnapshot.withdrawable2);
+
+    // move 5 weeks forward on L2
+    _moveDays({ day: 5 * 7, forward: true, isL2: true });
+
+    assertEq(locking.getWeek(), weekNoBefore + 5);
 
+    blocksTillNextWeekL2 = _calculateBlocksTillNextWeek({ isL2: true });
+
+    assertEq(blocksTillNextWeekL2, 5 * blocksTillNextWeekL1);
+    assertEq(locking.totalSupply(), afterSnapshot.totalSupply);
+    assertEq(locking.getPastTotalSupply(block.number - 3 * L2_WEEK), afterSnapshot.pastTotalSupply);
+    assertEq(locking.balanceOf(AIRDROP_CLAIMER_1), afterSnapshot.balance1);
+    assertEq(locking.balanceOf(AIRDROP_CLAIMER_2), afterSnapshot.balance2);
+    assertEq(locking.getVotes(AIRDROP_CLAIMER_1), afterSnapshot.votingPower1);
+    assertEq(locking.getVotes(AIRDROP_CLAIMER_2), afterSnapshot.votingPower2);
+    assertEq(locking.getPastVotes(AIRDROP_CLAIMER_1, block.number - 3 * L2_WEEK), afterSnapshot.pastVotingPower1);
+    assertEq(locking.getPastVotes(AIRDROP_CLAIMER_2, block.number - 3 * L2_WEEK), afterSnapshot.pastVotingPower2);
+    assertEq(locking.locked(AIRDROP_CLAIMER_1), afterSnapshot.lockedBalance1);
+    assertEq(locking.locked(AIRDROP_CLAIMER_2), afterSnapshot.lockedBalance2);
+    assertEq(locking.getAvailableForWithdraw(AIRDROP_CLAIMER_1), afterSnapshot.withdrawable1);
+    assertEq(locking.getAvailableForWithdraw(AIRDROP_CLAIMER_2), afterSnapshot.withdrawable2);
+
+    // move 5 days forward on L2
+    _moveDays({ day: 5, forward: true, isL2: true });
+    // we should be at the same week (TUE around 00:00)
     assertEq(locking.getWeek(), weekNoBefore + 5);
+    // move 1 day forward on L2 + 90 mins as buffer
+    _moveDays({ day: 1, forward: true, isL2: true });
+    vm.roll(block.number + 90 minutes);
+    // we should be at the next week (WED around 01:00)
+    assertEq(locking.getWeek(), weekNoBefore + 6);
+  }
+
+  // takes a snapshot of the locking contract at current block
+  function _takeSnapshot(address claimer1, address claimer2) internal view returns (LockingSnapshot memory snapshot) {
+    snapshot.totalSupply = locking.totalSupply();
+    snapshot.pastTotalSupply = locking.getPastTotalSupply(block.number - 3 * L1_WEEK);
+    snapshot.balance1 = locking.balanceOf(claimer1);
+    snapshot.balance2 = locking.balanceOf(claimer2);
+    snapshot.votingPower1 = locking.getVotes(claimer1);
+    snapshot.votingPower2 = locking.getVotes(claimer2);
+    snapshot.pastVotingPower1 = locking.getPastVotes(claimer1, block.number - 3 * L1_WEEK);
+    snapshot.pastVotingPower2 = locking.getPastVotes(claimer2, block.number - 3 * L1_WEEK);
+    snapshot.lockedBalance1 = locking.locked(claimer1);
+    snapshot.lockedBalance2 = locking.locked(claimer2);
+    snapshot.withdrawable1 = locking.getAvailableForWithdraw(claimer1);
+    snapshot.withdrawable2 = locking.getAvailableForWithdraw(claimer2);
+  }
+
+  // returns the number of blocks till the next week
+  // by calculating the first block of the next week and substracting the current block
+  function _calculateBlocksTillNextWeek(bool isL2) internal view returns (uint256) {
+    if (isL2) {
+      return ((L2_WEEK * (locking.getWeek() + uint256(locking.l2StartingPointWeek()) + 1)) + 507776) - block.number;
+    } else {
+      return ((L1_WEEK * (locking.getWeek() + locking.startingPointWeek() + 1)) + 89964) - block.number;
+    }
+  }
 
-    blocksTillNextWeek2 =
-      ((L2_WEEK * (locking.getWeek() + uint256(locking.l2StartingPointWeek()) + 1)) + 507776) -
-      block.number;
-
-    assertEq(blocksTillNextWeek2, 5 * blocksTillNextWeek);
-    assertEq(locking.totalSupply(), afterBalances.totalSupply);
-    assertEq(locking.getPastTotalSupply(block.number - 3 * L2_WEEK), afterBalances.pastTotalSupply);
-    assertEq(locking.balanceOf(AIRDROP_CLAIMER_1), afterBalances.balance1);
-    assertEq(locking.balanceOf(AIRDROP_CLAIMER_2), afterBalances.balance2);
-    assertEq(locking.getVotes(AIRDROP_CLAIMER_1), afterBalances.votingPower1);
-    assertEq(locking.getVotes(AIRDROP_CLAIMER_2), afterBalances.votingPower2);
-    assertEq(locking.getPastVotes(AIRDROP_CLAIMER_1, block.number - 3 * L2_WEEK), afterBalances.pastVotingPower1);
-    assertEq(locking.getPastVotes(AIRDROP_CLAIMER_2, block.number - 3 * L2_WEEK), afterBalances.pastVotingPower2);
-    assertEq(locking.locked(AIRDROP_CLAIMER_1), afterBalances.lockedBalance1);
-    assertEq(locking.locked(AIRDROP_CLAIMER_2), afterBalances.lockedBalance2);
-    assertEq(locking.getAvailableForWithdraw(AIRDROP_CLAIMER_1), afterBalances.withdrawable1);
-    assertEq(locking.getAvailableForWithdraw(AIRDROP_CLAIMER_2), afterBalances.withdrawable2);
+  // simulates the L2 upgrade by setting the necessary parameters
+  function _simulateL2Upgrade() internal {
+    vm.prank(mentoLabsMultisig);
+    locking.setL2TransitionBlock(block.number);
+    vm.prank(mentoLabsMultisig);
+    locking.setL2StartingPointWeek(20);
+    vm.prank(mentoLabsMultisig);
+    locking.setL2Shift(507776);
+    vm.prank(mentoLabsMultisig);
+    locking.setPaused(false);
   }
 
   // move days forward or backward on L1 or L2
-  function moveDays(uint256 day, bool forward, bool isL2) public {
+  function _moveDays(uint256 day, bool forward, bool isL2) internal {
     uint256 ts = vm.getBlockTimestamp();
     uint256 height = vm.getBlockNumber();
 
     uint256 newTs = forward ? ts + day * 1 days : ts - day * 1 days;
+
     uint256 blockChange = isL2 ? (day * 1 days) : ((day * 1 days) / 5);
     uint256 newHeight = forward ? height + blockChange : height - blockChange;
+
     vm.warp(newTs);
     vm.roll(newHeight);
-
-    ts = vm.getBlockTimestamp();
-    height = vm.getBlockNumber();
   }
 }

From 325c5fb3efd4d4f3fc10e50026719a0d703714bf Mon Sep 17 00:00:00 2001
From: baroooo <baranseltekin@gmail.com>
Date: Mon, 18 Nov 2024 15:32:26 +0100
Subject: [PATCH 04/16] chore: prettify

---
 contracts/governance/locking/LockingBase.sol  | 43 +++++++++----------
 test/fork/upgrades/LockingUpgradeForkTest.sol |  1 +
 .../governance/Locking/locking.upgrade.t.sol  | 37 +++++-----------
 test/utils/harnesses/LockingHarness.sol       |  3 +-
 4 files changed, 33 insertions(+), 51 deletions(-)

diff --git a/contracts/governance/locking/LockingBase.sol b/contracts/governance/locking/LockingBase.sol
index 0c6b3bdc..b8b1527e 100644
--- a/contracts/governance/locking/LockingBase.sol
+++ b/contracts/governance/locking/LockingBase.sol
@@ -93,7 +93,7 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
    */
   LibBrokenLine.BrokenLine public totalSupplyLine;
 
-  // *************** 
+  // ***************
   // New variables for L2 transition upgrade (4 slots)
   // ***************
   /**
@@ -105,7 +105,7 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
    */
   address public mentoLabsMultisig;
   /**
-   * @dev L2 starting point week number used to move the first week after the L2 transition to the last week before the L2 transition
+   * @dev L2 starting point week number
    */
   int256 public l2StartingPointWeek;
   /**
@@ -300,7 +300,7 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
    */
   function roundTimestamp(uint32 ts) public view returns (uint32) {
     require(!paused, "locking is paused");
-    
+
     if (ts < getEpochShift()) {
       return 0;
     }
@@ -316,7 +316,7 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
 
   /**
    * @notice method returns the amount of blocks to shift locking epoch to on L1 CELO.
-   * we move it to 00-00 UTC Wednesday (approx) by shifting 89964 blocks 
+   * we move it to 00-00 UTC Wednesday (approx) by shifting 89964 blocks
    */
   function getEpochShift() internal view virtual returns (uint32) {
     return 89964;
@@ -399,45 +399,42 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
     updateTotalSupplyLine(time);
   }
   /**
-    * @notice Sets the Mento Labs multisig address
-    * @param mentoLabsMultisig_ address of the Mento Labs multisig
-    *
-    */
+   * @notice Sets the Mento Labs multisig address
+   * @param mentoLabsMultisig_ address of the Mento Labs multisig
+   *
+   */
   function setMentoLabsMultisig(address mentoLabsMultisig_) external onlyOwner {
     mentoLabsMultisig = mentoLabsMultisig_;
   }
   /**
-    * @notice Sets the L2 transition block number and pauses locking and governance
-    * @param blockNo block number of the L2 transition
-    *
-    */
-  function setL2TransitionBlock(uint256 blockNo) external onlyMentoLabs{
-
+   * @notice Sets the L2 transition block number and pauses locking and governance
+   * @param blockNo block number of the L2 transition
+   *
+   */
+  function setL2TransitionBlock(uint256 blockNo) external onlyMentoLabs {
     l2Block = blockNo;
     paused = true;
   }
 
   /**
-    * @notice Sets the L2 shift amount
-    * @param l2Shift_ shift amount that will be used after L2 transition
+   * @notice Sets the L2 shift amount
+   * @param l2Shift_ shift amount that will be used after L2 transition
    */
-  function setL2Shift(uint32 l2Shift_) external onlyMentoLabs{
-
+  function setL2Shift(uint32 l2Shift_) external onlyMentoLabs {
     l2Shift = l2Shift_;
   }
 
   /**
-    * @notice Sets the L2 starting point week number
-    * @param l2StartingPointWeek_ starting point week number that will be used after L2 transition
+   * @notice Sets the L2 starting point week number
+   * @param l2StartingPointWeek_ starting point week number that will be used after L2 transition
    */
   function setL2StartingPointWeek(int256 l2StartingPointWeek_) external onlyMentoLabs {
-
     l2StartingPointWeek = l2StartingPointWeek_;
   }
 
   /**
-    * @notice Sets the paused flag
-    * @param paused_ flag to pause locking and governance
+   * @notice Sets the paused flag
+   * @param paused_ flag to pause locking and governance
    */
   function setPaused(bool paused_) external onlyMentoLabs {
     paused = paused_;
diff --git a/test/fork/upgrades/LockingUpgradeForkTest.sol b/test/fork/upgrades/LockingUpgradeForkTest.sol
index db6e8dfa..fcf5e08a 100644
--- a/test/fork/upgrades/LockingUpgradeForkTest.sol
+++ b/test/fork/upgrades/LockingUpgradeForkTest.sol
@@ -1,3 +1,4 @@
+/* solhint-disable max-line-length */
 // SPDX-License-Identifier: GPL-3.0-or-later
 pragma solidity ^0.8;
 
diff --git a/test/unit/governance/Locking/locking.upgrade.t.sol b/test/unit/governance/Locking/locking.upgrade.t.sol
index 9c71b423..2b42efaa 100644
--- a/test/unit/governance/Locking/locking.upgrade.t.sol
+++ b/test/unit/governance/Locking/locking.upgrade.t.sol
@@ -7,7 +7,6 @@ import { LockingTest } from "./LockingTest.sol";
 import "forge-std/console.sol";
 
 contract Upgrade_LockingTest is LockingTest {
-
   address public mentoLabs = makeAddr("MentoLabsMultisig");
 
   uint32 public l1Day;
@@ -15,7 +14,7 @@ contract Upgrade_LockingTest is LockingTest {
   uint32 public l1Week;
   uint32 public l2Week;
 
-   function setUp() public override{
+  function setUp() public override {
     super.setUp();
     l1Week = 7 days / 5;
     l2Week = 7 days;
@@ -40,7 +39,7 @@ contract Upgrade_LockingTest is LockingTest {
 
   function test_setMentoLabsMultisig_whenCalledByOwner_shouldSetMultisigAddress() public {
     assertEq(locking.mentoLabsMultisig(), address(0));
-    
+
     vm.prank(owner);
     locking.setMentoLabsMultisig(mentoLabs);
 
@@ -62,9 +61,9 @@ contract Upgrade_LockingTest is LockingTest {
   function test_setL2TransitionBlock_whenCalledByMentoMultisig_shouldSetL2BlockAndPause() public setMultisig {
     assertEq(locking.l2Block(), 0);
     assert(!locking.paused());
-    
+
     uint32 blockNumber = uint32(block.number + 100);
-    
+
     vm.prank(mentoLabs);
     locking.setL2TransitionBlock(blockNumber);
 
@@ -80,7 +79,7 @@ contract Upgrade_LockingTest is LockingTest {
 
   function test_setL2Shift_whenCalledByMentoMultisig_shouldSetL2BlockAndPause() public setMultisig {
     assertEq(locking.l2Shift(), 0);
-    
+
     vm.prank(mentoLabs);
     locking.setL2Shift(100);
 
@@ -95,8 +94,7 @@ contract Upgrade_LockingTest is LockingTest {
 
   function test_setL2StartingPointWeek_whenCalledByMentoMultisig_shouldSetL2BlockAndPause() public setMultisig {
     assertEq(locking.l2StartingPointWeek(), 0);
-    
-    
+
     vm.prank(mentoLabs);
     locking.setL2StartingPointWeek(100);
 
@@ -112,11 +110,10 @@ contract Upgrade_LockingTest is LockingTest {
     assert(!locking.paused());
 
     mentoToken.mint(alice, 1000000e18);
-    
+
     vm.prank(mentoLabs);
     locking.setPaused(true);
 
-
     assert(locking.paused());
 
     vm.expectRevert("locking is paused");
@@ -127,7 +124,6 @@ contract Upgrade_LockingTest is LockingTest {
     vm.prank(alice);
     locking.withdraw();
 
-
     vm.prank(mentoLabs);
     locking.setPaused(false);
 
@@ -153,7 +149,6 @@ contract Upgrade_LockingTest is LockingTest {
   }
 
   function test_getWeek_whenShiftAndStartingPointIs0_shouldReturnCorrectWeekNo() public l2LockingSetup(8, 0, 0) {
-
     // 2 + 8 weeks = 10 weeks on l1 = 2 weeks on l2
     assertEq(locking.getWeek(), 2);
     assertEq(locking.l2BlockTillNextPeriod(), l2Week);
@@ -189,7 +184,7 @@ contract Upgrade_LockingTest is LockingTest {
     assertEq(locking.getWeek(), l1WeekNo);
   }
 
-  function test_getWeek_whenL2StartingPointIsNegative_shouldReturnCorrectWeekNo() public l2LockingSetup(18,0,0) {
+  function test_getWeek_whenL2StartingPointIsNegative_shouldReturnCorrectWeekNo() public l2LockingSetup(18, 0, 0) {
     // l1 week no = 18 + 2 - 0 = 20
     uint32 l1WeekNo = 20;
     // l2 week no = (18 + 2) / 5 = 4
@@ -233,7 +228,6 @@ contract Upgrade_LockingTest is LockingTest {
     assertEq(locking.getWeek(), l1WeekNo);
     assertEq(locking.l2BlockTillNextPeriod(), l2Day);
 
-
     _incrementBlock(l2Day);
 
     assertEq(locking.getWeek(), l1WeekNo + 1);
@@ -252,24 +246,21 @@ contract Upgrade_LockingTest is LockingTest {
     // week no: 40
     _incrementBlock(l1Week * 20);
 
-
     uint256 totalSupplyL1W40 = locking.totalSupply();
-    uint256 pastTotalSupplyL1W30= locking.getPastTotalSupply(locking.blockNumberMocked() - l1Week * 10);
+    uint256 pastTotalSupplyL1W30 = locking.getPastTotalSupply(locking.blockNumberMocked() - l1Week * 10);
 
     // week no: 60
     _incrementBlock(l1Week * 20);
 
     uint256 totalSupplyL1W60 = locking.totalSupply();
-    uint256 pastTotalSupplyL1W50= locking.getPastTotalSupply(locking.blockNumberMocked() - l1Week * 10);
+    uint256 pastTotalSupplyL1W50 = locking.getPastTotalSupply(locking.blockNumberMocked() - l1Week * 10);
 
     // roll back to week 40
     _reduceBlock(l1Week * 20);
-    
 
     vm.prank(mentoLabs);
     locking.setL2TransitionBlock(l1Week * 40);
 
-
     vm.prank(mentoLabs);
     locking.setPaused(false);
 
@@ -298,7 +289,6 @@ contract Upgrade_LockingTest is LockingTest {
     // week no: 40
     _incrementBlock(l1Week * 20);
 
-
     uint256 balanceOfL1W40 = locking.balanceOf(alice);
     uint256 votesL1W40 = locking.getVotes(alice);
     uint256 pastVotesL1W30 = locking.getPastVotes(alice, locking.blockNumberMocked() - l1Week * 10);
@@ -312,12 +302,10 @@ contract Upgrade_LockingTest is LockingTest {
 
     // roll back to week 40
     _reduceBlock(l1Week * 20);
-    
 
     vm.prank(mentoLabs);
     locking.setL2TransitionBlock(l1Week * 40);
 
-
     vm.prank(mentoLabs);
     locking.setPaused(false);
 
@@ -348,7 +336,6 @@ contract Upgrade_LockingTest is LockingTest {
     // week no: 40
     _incrementBlock(l1Week * 20);
 
-
     uint256 lockedL1W40 = locking.locked(alice);
     uint256 withdrawableL1W40 = locking.getAvailableForWithdraw(alice);
 
@@ -356,11 +343,10 @@ contract Upgrade_LockingTest is LockingTest {
     _incrementBlock(l1Week * 20);
 
     uint256 lockedL1W60 = locking.locked(alice);
-    uint256 withdrawableL1W60 = locking.getAvailableForWithdraw(alice); 
+    uint256 withdrawableL1W60 = locking.getAvailableForWithdraw(alice);
 
     // roll back to week 40
     _reduceBlock(l1Week * 20);
-    
 
     vm.prank(mentoLabs);
     locking.setL2TransitionBlock(l1Week * 40);
@@ -380,5 +366,4 @@ contract Upgrade_LockingTest is LockingTest {
     assertEq(locking.locked(alice), lockedL1W60);
     assertEq(locking.getAvailableForWithdraw(alice), withdrawableL1W60);
   }
- 
 }
diff --git a/test/utils/harnesses/LockingHarness.sol b/test/utils/harnesses/LockingHarness.sol
index e4e3521c..5a1870fd 100644
--- a/test/utils/harnesses/LockingHarness.sol
+++ b/test/utils/harnesses/LockingHarness.sol
@@ -57,8 +57,7 @@ contract LockingHarness is Locking {
     return (L2_WEEK * uint256(int(currentWeek) + l2StartingPointWeek + 1)) + l2Shift - getBlockNumber();
   }
 
-
   function setStatingPointWeek(uint32 _week) external {
     startingPointWeek = _week;
-  } 
+  }
 }

From 5166a127536d785d6f3b81fa5e68c84d7e8f5942 Mon Sep 17 00:00:00 2001
From: baroooo <baranseltekin@gmail.com>
Date: Mon, 18 Nov 2024 15:53:38 +0100
Subject: [PATCH 05/16] feat: pack new vars

---
 contracts/governance/locking/LockingBase.sol  | 22 +++++++++----------
 test/fork/upgrades/LockingUpgradeForkTest.sol | 13 +++++------
 2 files changed, 16 insertions(+), 19 deletions(-)

diff --git a/contracts/governance/locking/LockingBase.sol b/contracts/governance/locking/LockingBase.sol
index b8b1527e..2a227d14 100644
--- a/contracts/governance/locking/LockingBase.sol
+++ b/contracts/governance/locking/LockingBase.sol
@@ -7,8 +7,6 @@ import "openzeppelin-contracts-upgradeable/contracts/access/OwnableUpgradeable.s
 import "openzeppelin-contracts-upgradeable/contracts/governance/utils/IVotesUpgradeable.sol";
 import "./libs/LibBrokenLine.sol";
 
-import "forge-std/console.sol";
-
 /**
  * @title LockingBase
  * @dev This abstract contract provides the foundational functionality
@@ -94,16 +92,12 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
   LibBrokenLine.BrokenLine public totalSupplyLine;
 
   // ***************
-  // New variables for L2 transition upgrade (4 slots)
+  // New variables for L2 transition upgrade (3 slots)
   // ***************
   /**
    * @dev L2 transtion block number
    */
   uint256 public l2Block;
-  /**
-   * @dev Address of the Mento Labs multisig
-   */
-  address public mentoLabsMultisig;
   /**
    * @dev L2 starting point week number
    */
@@ -112,6 +106,10 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
    * @dev Shift amount used after L2 transition to move the start of the epoch to 00-00 UTC Wednesday (approx)
    */
   uint32 public l2Shift;
+  /**
+   * @dev Address of the Mento Labs multisig
+   */
+  address public mentoLabsMultisig;
   /**
    * @dev Flag to pause locking and governance
    */
@@ -315,7 +313,7 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
   }
 
   /**
-   * @notice method returns the amount of blocks to shift locking epoch to on L1 CELO.
+   * @notice method returns the amount of blocks to shift locking epoch on L1 CELO.
    * we move it to 00-00 UTC Wednesday (approx) by shifting 89964 blocks
    */
   function getEpochShift() internal view virtual returns (uint32) {
@@ -408,11 +406,11 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
   }
   /**
    * @notice Sets the L2 transition block number and pauses locking and governance
-   * @param blockNo block number of the L2 transition
+   * @param l2Block_ block number of the L2 transition
    *
    */
-  function setL2TransitionBlock(uint256 blockNo) external onlyMentoLabs {
-    l2Block = blockNo;
+  function setL2TransitionBlock(uint256 l2Block_) external onlyMentoLabs {
+    l2Block = l2Block_;
     paused = true;
   }
 
@@ -440,5 +438,5 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
     paused = paused_;
   }
 
-  uint256[46] private __gap;
+  uint256[47] private __gap;
 }
diff --git a/test/fork/upgrades/LockingUpgradeForkTest.sol b/test/fork/upgrades/LockingUpgradeForkTest.sol
index fcf5e08a..7b1ffb21 100644
--- a/test/fork/upgrades/LockingUpgradeForkTest.sol
+++ b/test/fork/upgrades/LockingUpgradeForkTest.sol
@@ -10,6 +10,7 @@ import { ITransparentUpgradeableProxy } from "openzeppelin-contracts-next/contra
 
 // used to avoid stack too deep error
 struct LockingSnapshot {
+  uint256 weekNo;
   uint256 totalSupply;
   uint256 pastTotalSupply;
   uint256 balance1;
@@ -67,8 +68,6 @@ contract LockingUpgradeForkTest is BaseForkTest {
     // move 3 weeks forward on L1
     _moveDays({ day: 3 * 7, forward: true, isL2: false });
 
-    uint256 weekNoBefore = locking.getWeek();
-
     // Take snapshot 3 weeks after Nov 07
     beforeSnapshot = _takeSnapshot(AIRDROP_CLAIMER_1, AIRDROP_CLAIMER_2);
 
@@ -90,7 +89,7 @@ contract LockingUpgradeForkTest is BaseForkTest {
     // if the shift number is correct, the number of blocks till the next week should be 5 times the previous number
     assertEq(blocksTillNextWeekL2, 5 * blocksTillNextWeekL1);
 
-    assertEq(locking.getWeek(), weekNoBefore);
+    assertEq(locking.getWeek(), beforeSnapshot.weekNo);
     assertEq(locking.totalSupply(), beforeSnapshot.totalSupply);
     // the past values should be calculated using the L1 week value
     assertEq(locking.getPastTotalSupply(block.number - 3 * L1_WEEK), beforeSnapshot.pastTotalSupply);
@@ -108,11 +107,10 @@ contract LockingUpgradeForkTest is BaseForkTest {
     // move 5 weeks forward on L2
     _moveDays({ day: 5 * 7, forward: true, isL2: true });
 
-    assertEq(locking.getWeek(), weekNoBefore + 5);
-
     blocksTillNextWeekL2 = _calculateBlocksTillNextWeek({ isL2: true });
 
     assertEq(blocksTillNextWeekL2, 5 * blocksTillNextWeekL1);
+    assertEq(locking.getWeek(), afterSnapshot.weekNo);
     assertEq(locking.totalSupply(), afterSnapshot.totalSupply);
     assertEq(locking.getPastTotalSupply(block.number - 3 * L2_WEEK), afterSnapshot.pastTotalSupply);
     assertEq(locking.balanceOf(AIRDROP_CLAIMER_1), afterSnapshot.balance1);
@@ -129,16 +127,17 @@ contract LockingUpgradeForkTest is BaseForkTest {
     // move 5 days forward on L2
     _moveDays({ day: 5, forward: true, isL2: true });
     // we should be at the same week (TUE around 00:00)
-    assertEq(locking.getWeek(), weekNoBefore + 5);
+    assertEq(locking.getWeek(), afterSnapshot.weekNo);
     // move 1 day forward on L2 + 90 mins as buffer
     _moveDays({ day: 1, forward: true, isL2: true });
     vm.roll(block.number + 90 minutes);
     // we should be at the next week (WED around 01:00)
-    assertEq(locking.getWeek(), weekNoBefore + 6);
+    assertEq(locking.getWeek(), afterSnapshot.weekNo + 1);
   }
 
   // takes a snapshot of the locking contract at current block
   function _takeSnapshot(address claimer1, address claimer2) internal view returns (LockingSnapshot memory snapshot) {
+    snapshot.weekNo = locking.getWeek();
     snapshot.totalSupply = locking.totalSupply();
     snapshot.pastTotalSupply = locking.getPastTotalSupply(block.number - 3 * L1_WEEK);
     snapshot.balance1 = locking.balanceOf(claimer1);

From d8762c22005009c56a67d125d546f82485560cca Mon Sep 17 00:00:00 2001
From: baroooo <baranseltekin@gmail.com>
Date: Mon, 18 Nov 2024 16:23:53 +0100
Subject: [PATCH 06/16] chore: clean up unit tests

---
 test/fork/upgrades/LockingUpgradeForkTest.sol |  2 +-
 .../governance/Locking/locking.upgrade.t.sol  | 29 +++++++------------
 2 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/test/fork/upgrades/LockingUpgradeForkTest.sol b/test/fork/upgrades/LockingUpgradeForkTest.sol
index 7b1ffb21..4a996d32 100644
--- a/test/fork/upgrades/LockingUpgradeForkTest.sol
+++ b/test/fork/upgrades/LockingUpgradeForkTest.sol
@@ -131,7 +131,7 @@ contract LockingUpgradeForkTest is BaseForkTest {
     // move 1 day forward on L2 + 90 mins as buffer
     _moveDays({ day: 1, forward: true, isL2: true });
     vm.roll(block.number + 90 minutes);
-    // we should be at the next week (WED around 01:00)
+    // we should be at the next week (WED around 01:30)
     assertEq(locking.getWeek(), afterSnapshot.weekNo + 1);
   }
 
diff --git a/test/unit/governance/Locking/locking.upgrade.t.sol b/test/unit/governance/Locking/locking.upgrade.t.sol
index 2b42efaa..267e9cfe 100644
--- a/test/unit/governance/Locking/locking.upgrade.t.sol
+++ b/test/unit/governance/Locking/locking.upgrade.t.sol
@@ -4,8 +4,6 @@ pragma solidity 0.8.18;
 
 import { LockingTest } from "./LockingTest.sol";
 
-import "forge-std/console.sol";
-
 contract Upgrade_LockingTest is LockingTest {
   address public mentoLabs = makeAddr("MentoLabsMultisig");
 
@@ -23,7 +21,7 @@ contract Upgrade_LockingTest is LockingTest {
   }
 
   function test_initialSetup_shouldHaveCorrectValues() public view {
-    assertEq(locking.L2_WEEK(), 7 days);
+    assertEq(locking.L2_WEEK(), l2Week);
     assertEq(locking.mentoLabsMultisig(), address(0));
     assertEq(locking.l2Block(), 0);
     assertEq(locking.l2StartingPointWeek(), 0);
@@ -38,8 +36,6 @@ contract Upgrade_LockingTest is LockingTest {
   }
 
   function test_setMentoLabsMultisig_whenCalledByOwner_shouldSetMultisigAddress() public {
-    assertEq(locking.mentoLabsMultisig(), address(0));
-
     vm.prank(owner);
     locking.setMentoLabsMultisig(mentoLabs);
 
@@ -59,9 +55,6 @@ contract Upgrade_LockingTest is LockingTest {
   }
 
   function test_setL2TransitionBlock_whenCalledByMentoMultisig_shouldSetL2BlockAndPause() public setMultisig {
-    assertEq(locking.l2Block(), 0);
-    assert(!locking.paused());
-
     uint32 blockNumber = uint32(block.number + 100);
 
     vm.prank(mentoLabs);
@@ -78,8 +71,6 @@ contract Upgrade_LockingTest is LockingTest {
   }
 
   function test_setL2Shift_whenCalledByMentoMultisig_shouldSetL2BlockAndPause() public setMultisig {
-    assertEq(locking.l2Shift(), 0);
-
     vm.prank(mentoLabs);
     locking.setL2Shift(100);
 
@@ -93,8 +84,6 @@ contract Upgrade_LockingTest is LockingTest {
   }
 
   function test_setL2StartingPointWeek_whenCalledByMentoMultisig_shouldSetL2BlockAndPause() public setMultisig {
-    assertEq(locking.l2StartingPointWeek(), 0);
-
     vm.prank(mentoLabs);
     locking.setL2StartingPointWeek(100);
 
@@ -107,8 +96,6 @@ contract Upgrade_LockingTest is LockingTest {
   }
 
   function test_setPaused_whenCalledByMentoMultisig_shouldPauseContracts() public setMultisig {
-    assert(!locking.paused());
-
     mentoToken.mint(alice, 1000000e18);
 
     vm.prank(mentoLabs);
@@ -128,16 +115,22 @@ contract Upgrade_LockingTest is LockingTest {
     locking.setPaused(false);
 
     assert(!locking.paused());
+
+    vm.prank(alice);
+    locking.lock(alice, bob, 1000e18, 5, 5);
+
+    vm.prank(alice);
+    locking.withdraw();
   }
 
-  modifier l2LockingSetup(uint32 advanceWeeks_, uint32 startingPointWeek_, uint32 l1Shift_) {
+  modifier l2LockingSetup(uint32 advanceWeeks, uint32 startingPointWeek, uint32 l1Shift) {
     vm.prank(owner);
     locking.setMentoLabsMultisig(mentoLabs);
 
-    _incrementBlock(l1Week * advanceWeeks_);
+    _incrementBlock(l1Week * advanceWeeks);
 
-    locking.setStatingPointWeek(startingPointWeek_);
-    locking.setEpochShift(l1Shift_);
+    locking.setStatingPointWeek(startingPointWeek);
+    locking.setEpochShift(l1Shift);
 
     vm.prank(mentoLabs);
     locking.setL2TransitionBlock(block.number);

From 5872e8d2f8b79a6ebf14d624cc62fa322b8ea617 Mon Sep 17 00:00:00 2001
From: baroooo <baranseltekin@gmail.com>
Date: Mon, 18 Nov 2024 17:24:08 +0100
Subject: [PATCH 07/16] feat: events for setters

---
 contracts/governance/locking/LockingBase.sol | 29 ++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/contracts/governance/locking/LockingBase.sol b/contracts/governance/locking/LockingBase.sol
index 2a227d14..e409224c 100644
--- a/contracts/governance/locking/LockingBase.sol
+++ b/contracts/governance/locking/LockingBase.sol
@@ -156,6 +156,26 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
    * @dev set newMinSlopePeriod
    */
   event SetMinSlopePeriod(uint256 indexed newMinSlopePeriod);
+  /**
+   * @dev set new Mento Labs multisig address
+   */
+  event SetMentoLabsMultisig(address indexed mentoLabsMultisig);
+  /**
+   * @dev set new L2 transition block number
+   */
+  event SetL2TransitionBlock(uint256 indexed l2Block);
+  /**
+   * @dev set new L2 shift amount
+   */
+  event SetL2Shift(uint32 indexed l2Shift);
+  /**
+   * @dev set new L2 starting point week number
+   */
+  event SetL2StartingPointWeek(int256 indexed l2StartingPointWeek);
+  /**
+   * @dev set new paused flag
+   */
+  event SetPaused(bool indexed paused);
 
   /**
    * @dev Initializes the contract with token, starting point week, and minimum cliff and slope periods.
@@ -403,6 +423,7 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
    */
   function setMentoLabsMultisig(address mentoLabsMultisig_) external onlyOwner {
     mentoLabsMultisig = mentoLabsMultisig_;
+    emit SetMentoLabsMultisig(mentoLabsMultisig_);
   }
   /**
    * @notice Sets the L2 transition block number and pauses locking and governance
@@ -412,6 +433,8 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
   function setL2TransitionBlock(uint256 l2Block_) external onlyMentoLabs {
     l2Block = l2Block_;
     paused = true;
+
+    emit SetL2TransitionBlock(l2Block_);
   }
 
   /**
@@ -420,6 +443,8 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
    */
   function setL2Shift(uint32 l2Shift_) external onlyMentoLabs {
     l2Shift = l2Shift_;
+
+    emit SetL2Shift(l2Shift_);
   }
 
   /**
@@ -428,6 +453,8 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
    */
   function setL2StartingPointWeek(int256 l2StartingPointWeek_) external onlyMentoLabs {
     l2StartingPointWeek = l2StartingPointWeek_;
+
+    emit SetL2StartingPointWeek(l2StartingPointWeek_);
   }
 
   /**
@@ -436,6 +463,8 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
    */
   function setPaused(bool paused_) external onlyMentoLabs {
     paused = paused_;
+
+    emit SetPaused(paused_);
   }
 
   uint256[47] private __gap;

From b3aa923fa711f45706d658137288a74c5f25a3a4 Mon Sep 17 00:00:00 2001
From: baroooo <baranseltekin@gmail.com>
Date: Tue, 19 Nov 2024 11:10:57 +0100
Subject: [PATCH 08/16] fix: casting in epoch block calculations

---
 test/fork/upgrades/LockingUpgradeForkTest.sol | 4 ++--
 test/utils/harnesses/LockingHarness.sol       | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/test/fork/upgrades/LockingUpgradeForkTest.sol b/test/fork/upgrades/LockingUpgradeForkTest.sol
index 4a996d32..6a1914f5 100644
--- a/test/fork/upgrades/LockingUpgradeForkTest.sol
+++ b/test/fork/upgrades/LockingUpgradeForkTest.sol
@@ -156,9 +156,9 @@ contract LockingUpgradeForkTest is BaseForkTest {
   // by calculating the first block of the next week and substracting the current block
   function _calculateBlocksTillNextWeek(bool isL2) internal view returns (uint256) {
     if (isL2) {
-      return ((L2_WEEK * (locking.getWeek() + uint256(locking.l2StartingPointWeek()) + 1)) + 507776) - block.number;
+      return L2_WEEK * uint256(int256(locking.getWeek()) + locking.l2StartingPointWeek() + 1) + 507776 - block.number;
     } else {
-      return ((L1_WEEK * (locking.getWeek() + locking.startingPointWeek() + 1)) + 89964) - block.number;
+      return L1_WEEK * (locking.getWeek() + locking.startingPointWeek() + 1) + 89964 - block.number;
     }
   }
 
diff --git a/test/utils/harnesses/LockingHarness.sol b/test/utils/harnesses/LockingHarness.sol
index 5a1870fd..36b3cf56 100644
--- a/test/utils/harnesses/LockingHarness.sol
+++ b/test/utils/harnesses/LockingHarness.sol
@@ -54,7 +54,7 @@ contract LockingHarness is Locking {
 
   function l2BlockTillNextPeriod() external view returns (uint256) {
     uint256 currentWeek = this.getWeek();
-    return (L2_WEEK * uint256(int(currentWeek) + l2StartingPointWeek + 1)) + l2Shift - getBlockNumber();
+    return (L2_WEEK * uint256(int256(currentWeek) + l2StartingPointWeek + 1)) + l2Shift - getBlockNumber();
   }
 
   function setStatingPointWeek(uint32 _week) external {

From 3c82cec5f1bc4ad514a7f8298a98e84e5d69bc53 Mon Sep 17 00:00:00 2001
From: baroooo <baranseltekin@gmail.com>
Date: Thu, 21 Nov 2024 16:40:18 +0100
Subject: [PATCH 09/16] test: governance functions after l2 transition

---
 test/fork/upgrades/LockingUpgradeForkTest.sol | 98 ++++++++++++++++++-
 1 file changed, 94 insertions(+), 4 deletions(-)

diff --git a/test/fork/upgrades/LockingUpgradeForkTest.sol b/test/fork/upgrades/LockingUpgradeForkTest.sol
index 6a1914f5..f5c96f77 100644
--- a/test/fork/upgrades/LockingUpgradeForkTest.sol
+++ b/test/fork/upgrades/LockingUpgradeForkTest.sol
@@ -5,6 +5,8 @@ pragma solidity ^0.8;
 import { BaseForkTest } from "../BaseForkTest.sol";
 import { Locking } from "contracts/governance/locking/Locking.sol";
 import { GovernanceFactory } from "contracts/governance/GovernanceFactory.sol";
+import { MentoGovernor } from "contracts/governance/MentoGovernor.sol";
+import { MentoToken } from "contracts/governance/MentoToken.sol";
 import { ProxyAdmin } from "openzeppelin-contracts-next/contracts/proxy/transparent/ProxyAdmin.sol";
 import { ITransparentUpgradeableProxy } from "openzeppelin-contracts-next/contracts/proxy/transparent/TransparentUpgradeableProxy.sol";
 
@@ -36,6 +38,9 @@ contract LockingUpgradeForkTest is BaseForkTest {
   GovernanceFactory public governanceFactory = GovernanceFactory(0xee6CE2dbe788dFC38b8F583Da86cB9caf2C8cF5A);
   ProxyAdmin public proxyAdmin;
   Locking public locking;
+  MentoGovernor public mentoGovernor;
+  MentoToken public mentoToken;
+
   address public timelockController;
   address public newLockingImplementation;
 
@@ -48,6 +53,8 @@ contract LockingUpgradeForkTest is BaseForkTest {
     proxyAdmin = governanceFactory.proxyAdmin();
     locking = governanceFactory.locking();
     timelockController = address(governanceFactory.governanceTimelock());
+    mentoGovernor = governanceFactory.mentoGovernor();
+    mentoToken = governanceFactory.mentoToken();
 
     newLockingImplementation = address(new Locking());
     vm.prank(timelockController);
@@ -55,16 +62,16 @@ contract LockingUpgradeForkTest is BaseForkTest {
 
     vm.prank(timelockController);
     locking.setMentoLabsMultisig(mentoLabsMultisig);
-
-    // THU Nov-07-2024 00:00:23 +UTC
-    vm.roll(28653031);
-    vm.warp(1730937623);
   }
 
   function test_blockNoDependentCalculations_afterL2Transition_shouldWorkAsBefore() public {
     LockingSnapshot memory beforeSnapshot;
     LockingSnapshot memory afterSnapshot;
 
+    // THU Nov-07-2024 00:00:23 +UTC
+    vm.roll(28653031);
+    vm.warp(1730937623);
+
     // move 3 weeks forward on L1
     _moveDays({ day: 3 * 7, forward: true, isL2: false });
 
@@ -135,6 +142,89 @@ contract LockingUpgradeForkTest is BaseForkTest {
     assertEq(locking.getWeek(), afterSnapshot.weekNo + 1);
   }
 
+  function test_setPaused_shouldPauseGovernance() public {
+    _lockTokensForGovernance(AIRDROP_CLAIMER_1, 10_000_000e18);
+
+    vm.prank(mentoLabsMultisig);
+    locking.setPaused(true);
+
+    vm.prank(AIRDROP_CLAIMER_1);
+    vm.expectRevert("locking is paused");
+    mentoGovernor.propose(new address[](1), new uint256[](1), new bytes[](1), "Test proposal");
+
+    vm.prank(mentoLabsMultisig);
+    locking.setPaused(false);
+
+    vm.prank(AIRDROP_CLAIMER_1);
+    uint256 proposalId = mentoGovernor.propose(new address[](1), new uint256[](1), new bytes[](1), "Test proposal");
+
+    _moveDays(1, true, false);
+
+    vm.prank(mentoLabsMultisig);
+    locking.setPaused(true);
+
+    vm.prank(AIRDROP_CLAIMER_1);
+    vm.expectRevert("locking is paused");
+    mentoGovernor.castVote(proposalId, 1);
+  }
+
+  function test_governance_afterL2Transition_shouldWorkAsBefore() public {
+    _simulateL2Upgrade();
+
+    _moveDays(7, true, true);
+
+    uint256 votingPower1 = locking.getVotes(AIRDROP_CLAIMER_1);
+    uint256 votingPower2 = locking.getVotes(AIRDROP_CLAIMER_2);
+
+    uint256 lockId = _lockTokensForGovernance(AIRDROP_CLAIMER_1, 1_000_000e18);
+
+    assertEq(locking.getVotes(AIRDROP_CLAIMER_1), votingPower1 + 1_000_000e18);
+
+    vm.prank(AIRDROP_CLAIMER_1);
+    locking.delegateTo(lockId, AIRDROP_CLAIMER_2);
+
+    assertEq(locking.getVotes(AIRDROP_CLAIMER_1), votingPower1);
+    assertEq(locking.getVotes(AIRDROP_CLAIMER_2), votingPower2 + 1_000_000e18);
+
+    vm.prank(AIRDROP_CLAIMER_1);
+    locking.relock(lockId, AIRDROP_CLAIMER_1, 1_000_000e18, 1, 103);
+
+    assertEq(locking.getVotes(AIRDROP_CLAIMER_1), votingPower1 + 1_000_000e18);
+    assertEq(locking.getVotes(AIRDROP_CLAIMER_2), votingPower2);
+
+    vm.prank(AIRDROP_CLAIMER_1);
+    uint256 proposalId = mentoGovernor.propose(new address[](1), new uint256[](1), new bytes[](1), "Test proposal");
+
+    _moveDays(1, true, true);
+
+    vm.prank(AIRDROP_CLAIMER_1);
+    mentoGovernor.castVote(proposalId, 1);
+
+    vm.prank(AIRDROP_CLAIMER_2);
+    mentoGovernor.castVote(proposalId, 2);
+
+    _moveDays(5, true, true);
+
+    mentoGovernor.queue(proposalId);
+
+    _moveDays(2, true, true);
+
+    mentoGovernor.execute(proposalId);
+  }
+
+  // used to give locker enough power to be able to propose
+  function _lockTokensForGovernance(address locker, uint96 amount) internal returns (uint256 lockId) {
+    deal(address(mentoToken), locker, amount);
+
+    vm.prank(locker);
+    mentoToken.approve(address(locking), amount);
+
+    vm.prank(locker);
+    lockId = locking.lock(locker, locker, amount, 104, 0);
+
+    vm.roll(block.number + 1);
+  }
+
   // takes a snapshot of the locking contract at current block
   function _takeSnapshot(address claimer1, address claimer2) internal view returns (LockingSnapshot memory snapshot) {
     snapshot.weekNo = locking.getWeek();

From 981550deebc0554f3d08e7d45400bedd79fc4dad Mon Sep 17 00:00:00 2001
From: baroooo <baranseltekin@gmail.com>
Date: Wed, 27 Nov 2024 13:20:46 +0100
Subject: [PATCH 10/16] refactor: rename l2Block to l2TransitionBlock

---
 contracts/governance/locking/LockingBase.sol       | 14 +++++++-------
 test/unit/governance/Locking/locking.upgrade.t.sol | 10 +++++++---
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/contracts/governance/locking/LockingBase.sol b/contracts/governance/locking/LockingBase.sol
index e409224c..771dbaea 100644
--- a/contracts/governance/locking/LockingBase.sol
+++ b/contracts/governance/locking/LockingBase.sol
@@ -97,7 +97,7 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
   /**
    * @dev L2 transtion block number
    */
-  uint256 public l2Block;
+  uint256 public l2TransitionBlock;
   /**
    * @dev L2 starting point week number
    */
@@ -163,7 +163,7 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
   /**
    * @dev set new L2 transition block number
    */
-  event SetL2TransitionBlock(uint256 indexed l2Block);
+  event SetL2TransitionBlock(uint256 indexed l2TransitionBlock);
   /**
    * @dev set new L2 shift amount
    */
@@ -323,7 +323,7 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
       return 0;
     }
 
-    if (l2Block == 0 || ts < l2Block) {
+    if (l2TransitionBlock == 0 || ts < l2TransitionBlock) {
       uint32 shifted = ts - getEpochShift();
       return shifted / WEEK - uint32(startingPointWeek);
     } else {
@@ -427,14 +427,14 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
   }
   /**
    * @notice Sets the L2 transition block number and pauses locking and governance
-   * @param l2Block_ block number of the L2 transition
+   * @param l2TransitionBlock_ block number of the L2 transition
    *
    */
-  function setL2TransitionBlock(uint256 l2Block_) external onlyMentoLabs {
-    l2Block = l2Block_;
+  function setL2TransitionBlock(uint256 l2TransitionBlock_) external onlyMentoLabs {
+    l2TransitionBlock = l2TransitionBlock_;
     paused = true;
 
-    emit SetL2TransitionBlock(l2Block_);
+    emit SetL2TransitionBlock(l2TransitionBlock_);
   }
 
   /**
diff --git a/test/unit/governance/Locking/locking.upgrade.t.sol b/test/unit/governance/Locking/locking.upgrade.t.sol
index 267e9cfe..c7a99a1d 100644
--- a/test/unit/governance/Locking/locking.upgrade.t.sol
+++ b/test/unit/governance/Locking/locking.upgrade.t.sol
@@ -23,7 +23,7 @@ contract Upgrade_LockingTest is LockingTest {
   function test_initialSetup_shouldHaveCorrectValues() public view {
     assertEq(locking.L2_WEEK(), l2Week);
     assertEq(locking.mentoLabsMultisig(), address(0));
-    assertEq(locking.l2Block(), 0);
+    assertEq(locking.l2TransitionBlock(), 0);
     assertEq(locking.l2StartingPointWeek(), 0);
     assertEq(locking.l2Shift(), 0);
     assert(!locking.paused());
@@ -60,7 +60,7 @@ contract Upgrade_LockingTest is LockingTest {
     vm.prank(mentoLabs);
     locking.setL2TransitionBlock(blockNumber);
 
-    assertEq(locking.l2Block(), blockNumber);
+    assertEq(locking.l2TransitionBlock(), blockNumber);
     assert(locking.paused());
   }
 
@@ -123,7 +123,11 @@ contract Upgrade_LockingTest is LockingTest {
     locking.withdraw();
   }
 
-  modifier l2LockingSetup(uint32 advanceWeeks, uint32 startingPointWeek, uint32 l1Shift) {
+  modifier l2LockingSetup(
+    uint32 advanceWeeks,
+    uint32 startingPointWeek,
+    uint32 l1Shift
+  ) {
     vm.prank(owner);
     locking.setMentoLabsMultisig(mentoLabs);
 

From b52ea7ba5953ed00ad2e2b823a52763951ee2907 Mon Sep 17 00:00:00 2001
From: baroooo <baranseltekin@gmail.com>
Date: Wed, 27 Nov 2024 13:30:09 +0100
Subject: [PATCH 11/16] refactor: rename timestamp to blockNo

---
 contracts/governance/locking/Locking.sol       | 12 ++++++------
 contracts/governance/locking/LockingBase.sol   | 16 ++++++++--------
 contracts/governance/locking/LockingRelock.sol |  2 +-
 contracts/governance/locking/LockingVotes.sol  |  6 +++---
 4 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/contracts/governance/locking/Locking.sol b/contracts/governance/locking/Locking.sol
index 51e169b5..412b5f06 100644
--- a/contracts/governance/locking/Locking.sol
+++ b/contracts/governance/locking/Locking.sol
@@ -60,7 +60,7 @@ contract Locking is ILocking, LockingBase, LockingRelock, LockingVotes {
     counter++;
 
     uint32 currentBlock = getBlockNumber();
-    uint32 time = roundTimestamp(currentBlock);
+    uint32 time = getWeekNumber(currentBlock);
     addLines(account, _delegate, amount, slopePeriod, cliff, time, currentBlock);
     accounts[account].amount = accounts[account].amount + (amount);
 
@@ -92,7 +92,7 @@ contract Locking is ILocking, LockingBase, LockingRelock, LockingVotes {
   function getAvailableForWithdraw(address account) public view returns (uint96) {
     uint96 value = accounts[account].amount;
     uint32 currentBlock = getBlockNumber();
-    uint32 time = roundTimestamp(currentBlock);
+    uint32 time = getWeekNumber(currentBlock);
     uint96 bias = accounts[account].locked.actualValue(time, currentBlock);
     value = value - (bias);
     return value;
@@ -124,7 +124,7 @@ contract Locking is ILocking, LockingBase, LockingRelock, LockingVotes {
    * since the starting point week. The starting point is set during the contract initialization.
    */
   function getWeek() external view returns (uint256) {
-    return roundTimestamp(getBlockNumber());
+    return getWeekNumber(getBlockNumber());
   }
 
   /**
@@ -139,7 +139,7 @@ contract Locking is ILocking, LockingBase, LockingRelock, LockingVotes {
     address account = verifyLockOwner(id);
     address _delegate = locks[id].delegate;
     uint32 currentBlock = getBlockNumber();
-    uint32 time = roundTimestamp(currentBlock);
+    uint32 time = getWeekNumber(currentBlock);
     accounts[_delegate].balance.update(time);
     (uint96 bias, uint96 slope, uint32 cliff) = accounts[_delegate].balance.remove(id, time, currentBlock);
     LibBrokenLine.Line memory line = LibBrokenLine.Line(time, bias, slope, cliff);
@@ -158,7 +158,7 @@ contract Locking is ILocking, LockingBase, LockingRelock, LockingVotes {
       return 0;
     }
     uint32 currentBlock = getBlockNumber();
-    uint32 time = roundTimestamp(currentBlock);
+    uint32 time = getWeekNumber(currentBlock);
     return totalSupplyLine.actualValue(time, currentBlock);
   }
 
@@ -172,7 +172,7 @@ contract Locking is ILocking, LockingBase, LockingRelock, LockingVotes {
       return 0;
     }
     uint32 currentBlock = getBlockNumber();
-    uint32 time = roundTimestamp(currentBlock);
+    uint32 time = getWeekNumber(currentBlock);
     return accounts[account].balance.actualValue(time, currentBlock);
   }
 
diff --git a/contracts/governance/locking/LockingBase.sol b/contracts/governance/locking/LockingBase.sol
index 771dbaea..2f7a92bf 100644
--- a/contracts/governance/locking/LockingBase.sol
+++ b/contracts/governance/locking/LockingBase.sol
@@ -313,21 +313,21 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
   /**
    * @notice Calculates the week number for a given blocknumber
    * @dev It takes L2 transition into account to calculate the week number consistently
-   * @param ts block number
+   * @param blockNumber block number
    * @return week number the block number belongs to
    */
-  function roundTimestamp(uint32 ts) public view returns (uint32) {
+  function getWeekNumber(uint32 blockNumber) public view returns (uint32) {
     require(!paused, "locking is paused");
 
-    if (ts < getEpochShift()) {
+    if (blockNumber < getEpochShift()) {
       return 0;
     }
 
-    if (l2TransitionBlock == 0 || ts < l2TransitionBlock) {
-      uint32 shifted = ts - getEpochShift();
+    if (l2TransitionBlock == 0 || blockNumber < l2TransitionBlock) {
+      uint32 shifted = blockNumber - getEpochShift();
       return shifted / WEEK - uint32(startingPointWeek);
     } else {
-      uint32 shifted = ts - l2Shift;
+      uint32 shifted = blockNumber - l2Shift;
       return uint32(uint256(int256(uint256(shifted / L2_WEEK)) - l2StartingPointWeek));
     }
   }
@@ -404,7 +404,7 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
    * @param blockNumber block number until which to update lines
    */
   function updateAccountLinesBlockNumber(address account, uint32 blockNumber) external onlyOwner {
-    uint32 time = roundTimestamp(blockNumber);
+    uint32 time = getWeekNumber(blockNumber);
     updateAccountLines(account, time);
   }
 
@@ -413,7 +413,7 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
    * @param blockNumber block number until which to update line
    */
   function updateTotalSupplyLineBlockNumber(uint32 blockNumber) external onlyOwner {
-    uint32 time = roundTimestamp(blockNumber);
+    uint32 time = getWeekNumber(blockNumber);
     updateTotalSupplyLine(time);
   }
   /**
diff --git a/contracts/governance/locking/LockingRelock.sol b/contracts/governance/locking/LockingRelock.sol
index 089f07cb..f8fdaa36 100644
--- a/contracts/governance/locking/LockingRelock.sol
+++ b/contracts/governance/locking/LockingRelock.sol
@@ -30,7 +30,7 @@ abstract contract LockingRelock is LockingBase {
 
     address account = verifyLockOwner(id);
     uint32 currentBlock = getBlockNumber();
-    uint32 time = roundTimestamp(currentBlock);
+    uint32 time = getWeekNumber(currentBlock);
     verification(account, id, newAmount, newSlopePeriod, newCliff, time);
 
     address _delegate = locks[id].delegate;
diff --git a/contracts/governance/locking/LockingVotes.sol b/contracts/governance/locking/LockingVotes.sol
index 6d84417a..710e4845 100644
--- a/contracts/governance/locking/LockingVotes.sol
+++ b/contracts/governance/locking/LockingVotes.sol
@@ -15,7 +15,7 @@ contract LockingVotes is LockingBase {
    */
   function getVotes(address account) external view override returns (uint256) {
     uint32 currentBlock = getBlockNumber();
-    uint32 currentWeek = roundTimestamp(currentBlock);
+    uint32 currentWeek = getWeekNumber(currentBlock);
     return accounts[account].balance.actualValue(currentWeek, currentBlock);
   }
 
@@ -24,7 +24,7 @@ contract LockingVotes is LockingBase {
    * at the end of the last period
    */
   function getPastVotes(address account, uint256 blockNumber) external view override returns (uint256) {
-    uint32 currentWeek = roundTimestamp(uint32(blockNumber));
+    uint32 currentWeek = getWeekNumber(uint32(blockNumber));
     require(blockNumber < getBlockNumber() && currentWeek > 0, "block not yet mined");
 
     return accounts[account].balance.actualValue(currentWeek, uint32(blockNumber));
@@ -35,7 +35,7 @@ contract LockingVotes is LockingBase {
    * at the end of the last period
    */
   function getPastTotalSupply(uint256 blockNumber) external view override returns (uint256) {
-    uint32 currentWeek = roundTimestamp(uint32(blockNumber));
+    uint32 currentWeek = getWeekNumber(uint32(blockNumber));
     require(blockNumber < getBlockNumber() && currentWeek > 0, "block not yet mined");
 
     return totalSupplyLine.actualValue(currentWeek, uint32(blockNumber));

From 3941f85419dc420f0adc67a028cdbada17aed07c Mon Sep 17 00:00:00 2001
From: baroooo <baranseltekin@gmail.com>
Date: Wed, 27 Nov 2024 13:37:00 +0100
Subject: [PATCH 12/16] refactor: rename l2Shift to l2EpochShift

---
 contracts/governance/locking/LockingBase.sol     | 16 ++++++++--------
 test/fork/upgrades/LockingUpgradeForkTest.sol    |  2 +-
 .../governance/Locking/locking.upgrade.t.sol     | 16 ++++++++--------
 test/utils/harnesses/LockingHarness.sol          |  2 +-
 4 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/contracts/governance/locking/LockingBase.sol b/contracts/governance/locking/LockingBase.sol
index 2f7a92bf..ef06c2ef 100644
--- a/contracts/governance/locking/LockingBase.sol
+++ b/contracts/governance/locking/LockingBase.sol
@@ -105,7 +105,7 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
   /**
    * @dev Shift amount used after L2 transition to move the start of the epoch to 00-00 UTC Wednesday (approx)
    */
-  uint32 public l2Shift;
+  uint32 public l2EpochShift;
   /**
    * @dev Address of the Mento Labs multisig
    */
@@ -167,7 +167,7 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
   /**
    * @dev set new L2 shift amount
    */
-  event SetL2Shift(uint32 indexed l2Shift);
+  event SetL2EpochShift(uint32 indexed l2EpochShift);
   /**
    * @dev set new L2 starting point week number
    */
@@ -327,7 +327,7 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
       uint32 shifted = blockNumber - getEpochShift();
       return shifted / WEEK - uint32(startingPointWeek);
     } else {
-      uint32 shifted = blockNumber - l2Shift;
+      uint32 shifted = blockNumber - l2EpochShift;
       return uint32(uint256(int256(uint256(shifted / L2_WEEK)) - l2StartingPointWeek));
     }
   }
@@ -438,13 +438,13 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
   }
 
   /**
-   * @notice Sets the L2 shift amount
-   * @param l2Shift_ shift amount that will be used after L2 transition
+   * @notice Sets the L2 epoch shift amount
+   * @param l2EpochShift_ shift amount that will be used after L2 transition
    */
-  function setL2Shift(uint32 l2Shift_) external onlyMentoLabs {
-    l2Shift = l2Shift_;
+  function setL2EpochShift(uint32 l2EpochShift_) external onlyMentoLabs {
+    l2EpochShift = l2EpochShift_;
 
-    emit SetL2Shift(l2Shift_);
+    emit SetL2EpochShift(l2EpochShift_);
   }
 
   /**
diff --git a/test/fork/upgrades/LockingUpgradeForkTest.sol b/test/fork/upgrades/LockingUpgradeForkTest.sol
index f5c96f77..15a9690d 100644
--- a/test/fork/upgrades/LockingUpgradeForkTest.sol
+++ b/test/fork/upgrades/LockingUpgradeForkTest.sol
@@ -259,7 +259,7 @@ contract LockingUpgradeForkTest is BaseForkTest {
     vm.prank(mentoLabsMultisig);
     locking.setL2StartingPointWeek(20);
     vm.prank(mentoLabsMultisig);
-    locking.setL2Shift(507776);
+    locking.setL2EpochShift(507776);
     vm.prank(mentoLabsMultisig);
     locking.setPaused(false);
   }
diff --git a/test/unit/governance/Locking/locking.upgrade.t.sol b/test/unit/governance/Locking/locking.upgrade.t.sol
index c7a99a1d..92094c03 100644
--- a/test/unit/governance/Locking/locking.upgrade.t.sol
+++ b/test/unit/governance/Locking/locking.upgrade.t.sol
@@ -25,7 +25,7 @@ contract Upgrade_LockingTest is LockingTest {
     assertEq(locking.mentoLabsMultisig(), address(0));
     assertEq(locking.l2TransitionBlock(), 0);
     assertEq(locking.l2StartingPointWeek(), 0);
-    assertEq(locking.l2Shift(), 0);
+    assertEq(locking.l2EpochShift(), 0);
     assert(!locking.paused());
   }
 
@@ -64,17 +64,17 @@ contract Upgrade_LockingTest is LockingTest {
     assert(locking.paused());
   }
 
-  function test_setL2Shift_whenCalledByNonMentoMultisig_shouldRevert() public setMultisig {
+  function test_setL2EpochShift_whenCalledByNonMentoMultisig_shouldRevert() public setMultisig {
     vm.prank(alice);
     vm.expectRevert("caller is not MentoLabs multisig");
-    locking.setL2Shift(100);
+    locking.setL2EpochShift(100);
   }
 
-  function test_setL2Shift_whenCalledByMentoMultisig_shouldSetL2BlockAndPause() public setMultisig {
+  function test_setL2EpochShift_whenCalledByMentoMultisig_shouldSetL2BlockAndPause() public setMultisig {
     vm.prank(mentoLabs);
-    locking.setL2Shift(100);
+    locking.setL2EpochShift(100);
 
-    assertEq(locking.l2Shift(), 100);
+    assertEq(locking.l2EpochShift(), 100);
   }
 
   function test_setL2StartingPointWeek_whenCalledByNonMentoMultisig_shouldRevert() public setMultisig {
@@ -209,9 +209,9 @@ contract Upgrade_LockingTest is LockingTest {
     locking.setL2StartingPointWeek(-11);
 
     vm.prank(mentoLabs);
-    locking.setL2Shift(l2Day * 3);
+    locking.setL2EpochShift(l2Day * 3);
 
-    // after the L2 starting point week and l2Shift are set, the timing should be equal to the l1 timing
+    // after the L2 starting point week and l2EpochShift are set, the timing should be equal to the l1 timing
     assertEq(locking.getWeek(), l1WeekNo);
     assertEq(locking.l2BlockTillNextPeriod(), l2Day * 3);
 
diff --git a/test/utils/harnesses/LockingHarness.sol b/test/utils/harnesses/LockingHarness.sol
index 36b3cf56..ab995cc4 100644
--- a/test/utils/harnesses/LockingHarness.sol
+++ b/test/utils/harnesses/LockingHarness.sol
@@ -54,7 +54,7 @@ contract LockingHarness is Locking {
 
   function l2BlockTillNextPeriod() external view returns (uint256) {
     uint256 currentWeek = this.getWeek();
-    return (L2_WEEK * uint256(int256(currentWeek) + l2StartingPointWeek + 1)) + l2Shift - getBlockNumber();
+    return (L2_WEEK * uint256(int256(currentWeek) + l2StartingPointWeek + 1)) + l2EpochShift - getBlockNumber();
   }
 
   function setStatingPointWeek(uint32 _week) external {

From ee1040b2f8412856eb007d2e10e179b9a62ccb5e Mon Sep 17 00:00:00 2001
From: baroooo <baranseltekin@gmail.com>
Date: Wed, 27 Nov 2024 13:47:40 +0100
Subject: [PATCH 13/16] chore: prettier

---
 test/unit/governance/Locking/locking.upgrade.t.sol | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/test/unit/governance/Locking/locking.upgrade.t.sol b/test/unit/governance/Locking/locking.upgrade.t.sol
index 92094c03..6d74cbc6 100644
--- a/test/unit/governance/Locking/locking.upgrade.t.sol
+++ b/test/unit/governance/Locking/locking.upgrade.t.sol
@@ -123,11 +123,7 @@ contract Upgrade_LockingTest is LockingTest {
     locking.withdraw();
   }
 
-  modifier l2LockingSetup(
-    uint32 advanceWeeks,
-    uint32 startingPointWeek,
-    uint32 l1Shift
-  ) {
+  modifier l2LockingSetup(uint32 advanceWeeks, uint32 startingPointWeek, uint32 l1Shift) {
     vm.prank(owner);
     locking.setMentoLabsMultisig(mentoLabs);
 

From d025c750602ca6ae1f2252442165fdbbd2447d4a Mon Sep 17 00:00:00 2001
From: baroooo <baranseltekin@gmail.com>
Date: Thu, 28 Nov 2024 12:13:26 +0100
Subject: [PATCH 14/16] refactor: getEpochShift to work with l2Shift

---
 contracts/governance/locking/LockingBase.sol  | 32 +++++++++++++------
 .../governance/Locking/locking.upgrade.t.sol  | 22 ++++++-------
 test/utils/harnesses/LockingHarness.sol       | 20 +++++++-----
 3 files changed, 46 insertions(+), 28 deletions(-)

diff --git a/contracts/governance/locking/LockingBase.sol b/contracts/governance/locking/LockingBase.sol
index ef06c2ef..3ec18357 100644
--- a/contracts/governance/locking/LockingBase.sol
+++ b/contracts/governance/locking/LockingBase.sol
@@ -313,31 +313,45 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
   /**
    * @notice Calculates the week number for a given blocknumber
    * @dev It takes L2 transition into account to calculate the week number consistently
-   * @param blockNumber block number
+   * @param blockNumber block number to calculate the week number for
    * @return week number the block number belongs to
    */
   function getWeekNumber(uint32 blockNumber) public view returns (uint32) {
     require(!paused, "locking is paused");
 
-    if (blockNumber < getEpochShift()) {
+    if (blockNumber < _getEpochShift(blockNumber)) {
       return 0;
     }
+    uint32 shifted = blockNumber - _getEpochShift(blockNumber);
 
-    if (l2TransitionBlock == 0 || blockNumber < l2TransitionBlock) {
-      uint32 shifted = blockNumber - getEpochShift();
+    if (_isPreL2Transition(blockNumber)) {
       return shifted / WEEK - uint32(startingPointWeek);
     } else {
-      uint32 shifted = blockNumber - l2EpochShift;
       return uint32(uint256(int256(uint256(shifted / L2_WEEK)) - l2StartingPointWeek));
     }
   }
 
   /**
-   * @notice method returns the amount of blocks to shift locking epoch on L1 CELO.
-   * we move it to 00-00 UTC Wednesday (approx) by shifting 89964 blocks
+   * @notice Returns the epoch shift based on L2 transition status
+   * @dev Epoch shift is the amount of blocks to move the epoch start to 00-00 UTC Wednesday (approx).
+   * @dev l2EpochShift will be moved into a constant once L2 transition is complete.
+   * @param blockNumber block number to calculate the shift for
+   * @return shift amount in blocks (89964 for L1, l2EpochShift for L2)
    */
-  function getEpochShift() internal view virtual returns (uint32) {
-    return 89964;
+  function _getEpochShift(uint32 blockNumber) internal view virtual returns (uint32) {
+    if (_isPreL2Transition(blockNumber)) {
+      return 89964;
+    }
+    return l2EpochShift;
+  }
+
+  /**
+   * @notice Determines if a block is before the L2 transition point
+   * @param blockNumber block number to check
+   * @return true if before L2 transition, false if after
+   */
+  function _isPreL2Transition(uint32 blockNumber) internal view returns (bool) {
+    return l2TransitionBlock == 0 || blockNumber < l2TransitionBlock;
   }
 
   /**
diff --git a/test/unit/governance/Locking/locking.upgrade.t.sol b/test/unit/governance/Locking/locking.upgrade.t.sol
index 6d74cbc6..30a09bac 100644
--- a/test/unit/governance/Locking/locking.upgrade.t.sol
+++ b/test/unit/governance/Locking/locking.upgrade.t.sol
@@ -144,22 +144,22 @@ contract Upgrade_LockingTest is LockingTest {
   function test_getWeek_whenShiftAndStartingPointIs0_shouldReturnCorrectWeekNo() public l2LockingSetup(8, 0, 0) {
     // 2 + 8 weeks = 10 weeks on l1 = 2 weeks on l2
     assertEq(locking.getWeek(), 2);
-    assertEq(locking.l2BlockTillNextPeriod(), l2Week);
+    assertEq(locking.blockTillNextPeriod(), l2Week);
 
     _incrementBlock(l2Day * 3);
 
     assertEq(locking.getWeek(), 2);
-    assertEq(locking.l2BlockTillNextPeriod(), l2Day * 4);
+    assertEq(locking.blockTillNextPeriod(), l2Day * 4);
 
     _incrementBlock(l2Day * 5);
 
     assertEq(locking.getWeek(), 3);
-    assertEq(locking.l2BlockTillNextPeriod(), l2Day * 6);
+    assertEq(locking.blockTillNextPeriod(), l2Day * 6);
 
     _incrementBlock(l2Day * 8);
 
     assertEq(locking.getWeek(), 4);
-    assertEq(locking.l2BlockTillNextPeriod(), l2Day * 5);
+    assertEq(locking.blockTillNextPeriod(), l2Day * 5);
   }
 
   function test_getWeek_whenL2StartingPointIsPositive_shouldReturnCorrectWeekNo() public l2LockingSetup(198, 190, 0) {
@@ -167,7 +167,7 @@ contract Upgrade_LockingTest is LockingTest {
     uint32 l1WeekNo = 10;
     // l2 week no = (198 + 2) / 5 = 40
     assertEq(locking.getWeek(), 40);
-    assertEq(locking.l2BlockTillNextPeriod(), l2Week);
+    assertEq(locking.blockTillNextPeriod(), l2Week);
 
     // l2WeekNo - l1WeekNo = 40 - 10 = 30
     vm.prank(mentoLabs);
@@ -182,7 +182,7 @@ contract Upgrade_LockingTest is LockingTest {
     uint32 l1WeekNo = 20;
     // l2 week no = (18 + 2) / 5 = 4
     assertEq(locking.getWeek(), 4);
-    assertEq(locking.l2BlockTillNextPeriod(), l2Week);
+    assertEq(locking.blockTillNextPeriod(), l2Week);
 
     // l2WeekNo - l1WeekNo = 4 - 20 = -16
     vm.prank(mentoLabs);
@@ -198,7 +198,7 @@ contract Upgrade_LockingTest is LockingTest {
 
     // l2 week no = (18 + 2) / 5 = 4
     assertEq(locking.getWeek(), 4);
-    assertEq(locking.l2BlockTillNextPeriod(), l2Week);
+    assertEq(locking.blockTillNextPeriod(), l2Week);
 
     // l2WeekNo - l1WeekNo = 4 - 14 - 1 = -11
     vm.prank(mentoLabs);
@@ -209,22 +209,22 @@ contract Upgrade_LockingTest is LockingTest {
 
     // after the L2 starting point week and l2EpochShift are set, the timing should be equal to the l1 timing
     assertEq(locking.getWeek(), l1WeekNo);
-    assertEq(locking.l2BlockTillNextPeriod(), l2Day * 3);
+    assertEq(locking.blockTillNextPeriod(), l2Day * 3);
 
     _incrementBlock(l2Day);
 
     assertEq(locking.getWeek(), l1WeekNo);
-    assertEq(locking.l2BlockTillNextPeriod(), l2Day * 2);
+    assertEq(locking.blockTillNextPeriod(), l2Day * 2);
 
     _incrementBlock(l2Day);
 
     assertEq(locking.getWeek(), l1WeekNo);
-    assertEq(locking.l2BlockTillNextPeriod(), l2Day);
+    assertEq(locking.blockTillNextPeriod(), l2Day);
 
     _incrementBlock(l2Day);
 
     assertEq(locking.getWeek(), l1WeekNo + 1);
-    assertEq(locking.l2BlockTillNextPeriod(), l2Week);
+    assertEq(locking.blockTillNextPeriod(), l2Week);
   }
 
   function test_totalSupply_whenCalledAfterL2Transition_shouldReturnCorrectValues() public setMultisig {
diff --git a/test/utils/harnesses/LockingHarness.sol b/test/utils/harnesses/LockingHarness.sol
index ab995cc4..6b956ec7 100644
--- a/test/utils/harnesses/LockingHarness.sol
+++ b/test/utils/harnesses/LockingHarness.sol
@@ -35,8 +35,11 @@ contract LockingHarness is Locking {
     (lockAmount, lockSlope) = getLock(amount, slope, cliff);
   }
 
-  function getEpochShift() internal view override returns (uint32) {
-    return epochShift;
+  function _getEpochShift(uint32) internal view override returns (uint32) {
+    if (_isPreL2Transition(getBlockNumber())) {
+      return epochShift;
+    }
+    return l2EpochShift;
   }
 
   function setEpochShift(uint32 _epochShift) external {
@@ -49,12 +52,13 @@ contract LockingHarness is Locking {
 
   function blockTillNextPeriod() external view returns (uint256) {
     uint256 currentWeek = this.getWeek();
-    return (WEEK * (currentWeek + startingPointWeek + 1)) + getEpochShift() - getBlockNumber();
-  }
-
-  function l2BlockTillNextPeriod() external view returns (uint256) {
-    uint256 currentWeek = this.getWeek();
-    return (L2_WEEK * uint256(int256(currentWeek) + l2StartingPointWeek + 1)) + l2EpochShift - getBlockNumber();
+    if (_isPreL2Transition(getBlockNumber())) {
+      return (WEEK * (currentWeek + startingPointWeek + 1)) + _getEpochShift(getBlockNumber()) - getBlockNumber();
+    }
+    return
+      (L2_WEEK * uint256(int256(currentWeek) + l2StartingPointWeek + 1)) +
+      _getEpochShift(getBlockNumber()) -
+      getBlockNumber();
   }
 
   function setStatingPointWeek(uint32 _week) external {

From c3eb32d477d857c1da91f05d8304b0dbab847fb6 Mon Sep 17 00:00:00 2001
From: baroooo <baranseltekin@gmail.com>
Date: Thu, 28 Nov 2024 12:21:46 +0100
Subject: [PATCH 15/16] refactor: l1 epoch shift into a constant

---
 contracts/governance/locking/LockingBase.sol | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/contracts/governance/locking/LockingBase.sol b/contracts/governance/locking/LockingBase.sol
index 3ec18357..5a5ee0b0 100644
--- a/contracts/governance/locking/LockingBase.sol
+++ b/contracts/governance/locking/LockingBase.sol
@@ -20,12 +20,14 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
    * @dev Duration of a week in blocks on the CELO blockchain before the L2 transition (5 seconds per block)
    */
   uint32 public constant WEEK = 120_960;
-
   /**
    * @dev Duration of a week in blocks on the CELO blockchain after the L2 transition (1 seconds per block)
    */
   uint32 public constant L2_WEEK = 604_800;
-
+  /**
+   * @dev Epoch shift for L1
+   */
+  uint32 public constant L1_EPOCH_SHIFT = 89964;
   /**
    * @dev Maximum allowable cliff period for token locks in weeks
    */
@@ -334,13 +336,13 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
   /**
    * @notice Returns the epoch shift based on L2 transition status
    * @dev Epoch shift is the amount of blocks to move the epoch start to 00-00 UTC Wednesday (approx).
-   * @dev l2EpochShift will be moved into a constant once L2 transition is complete.
+   * @dev l2EpochShift will be moved into a constant once L2 transition is complete and stable.
    * @param blockNumber block number to calculate the shift for
-   * @return shift amount in blocks (89964 for L1, l2EpochShift for L2)
+   * @return shift amount in blocks (L1_EPOCH_SHIFT for L1, l2EpochShift for L2)
    */
   function _getEpochShift(uint32 blockNumber) internal view virtual returns (uint32) {
     if (_isPreL2Transition(blockNumber)) {
-      return 89964;
+      return L1_EPOCH_SHIFT;
     }
     return l2EpochShift;
   }

From 9cb45aec75192f2d2da542ec956ccf9d8dbb397a Mon Sep 17 00:00:00 2001
From: baroooo <baranseltekin@gmail.com>
Date: Mon, 2 Dec 2024 13:36:05 +0100
Subject: [PATCH 16/16] chore: fix comments and spacing

---
 contracts/governance/locking/LockingBase.sol | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/contracts/governance/locking/LockingBase.sol b/contracts/governance/locking/LockingBase.sol
index 5a5ee0b0..6953b29c 100644
--- a/contracts/governance/locking/LockingBase.sol
+++ b/contracts/governance/locking/LockingBase.sol
@@ -97,7 +97,7 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
   // New variables for L2 transition upgrade (3 slots)
   // ***************
   /**
-   * @dev L2 transtion block number
+   * @dev L2 transition block number
    */
   uint256 public l2TransitionBlock;
   /**
@@ -279,7 +279,7 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
    * @param amount of tokens to lock
    * @param slopePeriod period over which the tokens will unlock
    * @param cliff initial period during which tokens remain locked and do not start unlocking
-   **/
+   */
   function getLock(
     uint96 amount,
     uint32 slopePeriod,
@@ -432,19 +432,19 @@ abstract contract LockingBase is OwnableUpgradeable, IVotesUpgradeable {
     uint32 time = getWeekNumber(blockNumber);
     updateTotalSupplyLine(time);
   }
+
   /**
    * @notice Sets the Mento Labs multisig address
    * @param mentoLabsMultisig_ address of the Mento Labs multisig
-   *
    */
   function setMentoLabsMultisig(address mentoLabsMultisig_) external onlyOwner {
     mentoLabsMultisig = mentoLabsMultisig_;
     emit SetMentoLabsMultisig(mentoLabsMultisig_);
   }
+
   /**
    * @notice Sets the L2 transition block number and pauses locking and governance
    * @param l2TransitionBlock_ block number of the L2 transition
-   *
    */
   function setL2TransitionBlock(uint256 l2TransitionBlock_) external onlyMentoLabs {
     l2TransitionBlock = l2TransitionBlock_;