From 8f322269320c1db81fb16b4a98d961b67285c117 Mon Sep 17 00:00:00 2001
From: Richard Melkonian <r.v.melkonian@gmail.com>
Date: Fri, 10 May 2024 14:21:45 +0100
Subject: [PATCH] add epoch unit test

---
 protocol-units/dispute/src/MCR.sol    |  4 ++-
 protocol-units/dispute/test/MCR.t.sol | 36 +++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)
 create mode 100644 protocol-units/dispute/test/MCR.t.sol

diff --git a/protocol-units/dispute/src/MCR.sol b/protocol-units/dispute/src/MCR.sol
index 128157465..f275d06c3 100644
--- a/protocol-units/dispute/src/MCR.sol
+++ b/protocol-units/dispute/src/MCR.sol
@@ -55,6 +55,7 @@ contract MCR {
     uint256 public delta = 1 * SECONDS_IN_DAY;
     uint256 public p = 1 * SECONDS_IN_MINUTE;
     uint256 public supermajorityStake;
+    uint256 public epochStartTimestamp;
 
     mapping(address => Validator) public validators;
     mapping(bytes32 => Dispute) public disputes;
@@ -76,7 +77,7 @@ contract MCR {
     constructor(
         uint256 _delta,
         uint256 _supermajorityStake,
-        uint256 _epochDurationInDays,
+        uint256 _epochDurationInDays
     ) {
         delta = _delta;
         supermajorityStake = _supermajorityStake;
@@ -88,6 +89,7 @@ contract MCR {
         if (epochsPassed > 0) {
             currentEpoch += epochsPassed;
             epochStartTimestamp += epochsPassed * epochDuration;
+            epochStartTimestamp = block.timestamp;
         }
     }
 
diff --git a/protocol-units/dispute/test/MCR.t.sol b/protocol-units/dispute/test/MCR.t.sol
new file mode 100644
index 000000000..71d9cc405
--- /dev/null
+++ b/protocol-units/dispute/test/MCR.t.sol
@@ -0,0 +1,36 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.13;
+
+import "forge-std/Test.sol";
+import "../src/MCR.sol";
+
+contract MCRTest is Test {
+    MCR public mcr;
+    uint256 public epochDuration = 7 days;
+
+    function setUp() public {
+        mcr = new MCR(1 days, 100 ether, epochDuration / 1 days);
+    }
+
+    function testUpdateEpoch() public {
+        // Test initial epoch
+        assertEq(mcr.currentEpoch(), 0);
+        assertEq(mcr.epochStartTimestamp(), block.timestamp);
+
+        // Advance time by 3 epochs
+        vm.warp(block.timestamp + 3 * epochDuration);
+
+        // Call updateEpoch and check updated values
+        mcr.updateEpoch();
+        assertEq(mcr.currentEpoch(), 3);
+        assertEq(mcr.epochStartTimestamp(), block.timestamp);
+
+        // Advance time by 1 epoch and 1 day
+        vm.warp(block.timestamp + epochDuration + 1 days);
+
+        // Call updateEpoch and check updated values
+        mcr.updateEpoch();
+        assertEq(mcr.currentEpoch(), 4);
+        assertEq(mcr.epochStartTimestamp(), block.timestamp - 1 days);
+    }
+}
\ No newline at end of file