added more tests

Layr-Labs · Oct 3, 2024 · 6689328 · 6689328
1 parent 53ebbff
commit 6689328
Show file tree

Hide file tree

Showing 6 changed files with 114 additions and 26 deletions.
diff --git a/contracts/mocks/MockStrategy.sol b/contracts/mocks/MockStrategy.sol
@@ -0,0 +1,66 @@
+// SPDX-License-Identifier: BUSL-1.1
+pragma solidity ^0.8.0;
+
+import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "@eigenlayer/contracts/interfaces/IStrategy.sol";
+
+contract MockStrategy is IStrategy {
+    IERC20 public override underlyingToken;
+    uint256 public override totalShares;
+    mapping(address => uint256) public userShares;
+    uint256 public constant EXCHANGE_RATE = 1e18; // 1:1 exchange rate for simplicity
+
+    constructor(IERC20 _underlyingToken) {
+        underlyingToken = _underlyingToken;
+        emit StrategyTokenSet(_underlyingToken, 18); // Assuming 18 decimals for simplicity
+    }
+
+    function deposit(IERC20 token, uint256 amount) external override returns (uint256) {
+        require(token == underlyingToken, "Invalid token");
+        uint256 newShares = amount;
+        totalShares += newShares;
+        userShares[msg.sender] += newShares;
+        emit ExchangeRateEmitted(EXCHANGE_RATE);
+        return newShares;
+    }
+
+    function withdraw(address recipient, IERC20 token, uint256 amountShares) external override {
+        require(token == underlyingToken, "Invalid token");
+        require(userShares[msg.sender] >= amountShares, "Insufficient shares");
+        userShares[msg.sender] -= amountShares;
+        totalShares -= amountShares;
+        underlyingToken.transfer(recipient, amountShares);
+    }
+
+    function sharesToUnderlying(uint256 amountShares) external pure override returns (uint256) {
+        return amountShares;
+    }
+
+    function underlyingToShares(uint256 amountUnderlying) external pure override returns (uint256) {
+        return amountUnderlying;
+    }
+
+    function userUnderlying(address user) external view override returns (uint256) {
+        return userShares[user];
+    }
+
+    function shares(address user) external view override returns (uint256) {
+        return userShares[user];
+    }
+
+    function sharesToUnderlyingView(uint256 amountShares) external pure override returns (uint256) {
+        return amountShares;
+    }
+
+    function underlyingToSharesView(uint256 amountUnderlying) external pure override returns (uint256) {
+        return amountUnderlying;
+    }
+
+    function userUnderlyingView(address user) external view override returns (uint256) {
+        return userShares[user];
+    }
+
+    function explanation() external pure override returns (string memory) {
+        return "Mock Strategy for testing purposes";
+    }
+}
diff --git a/contracts/payments.json b/contracts/payments.json
@@ -0,0 +1,10 @@
+{
+  "leaves": [
+    "0x0000000000000000000000000000000000000000000000000000000000000001",
+    "0x0000000000000000000000000000000000000000000000000000000000000002"
+  ],
+  "tokenLeaves": [
+    "0x0000000000000000000000000000000000000000000000000000000000000003",
+    "0x0000000000000000000000000000000000000000000000000000000000000004"
+  ]
+}
diff --git a/contracts/script/utils/SetupPaymentsLib.sol b/contracts/script/utils/SetupPaymentsLib.sol
@@ -157,7 +157,7 @@ library SetupPaymentsLib {
         return abi.decode(data, (PaymentLeaves));
     }
 
-    function generateMerkleProof(bytes32[] memory leaves, uint256 index) internal pure returns (bytes memory) {
+    function generateMerkleProof(bytes32[] memory leaves, uint256 index) public pure returns (bytes memory) {
         require(leaves.length > 0, "Leaves array cannot be empty");
         require(index < leaves.length, "Index out of bounds");
 

diff --git a/contracts/test/SetupPaymentsLib.t.sol b/contracts/test/SetupPaymentsLib.t.sol
@@ -72,30 +72,46 @@ contract SetupPaymentsLibTest is Test, TestConstants {
     // }
 
 
-    // function testWriteLeavesToJson() public {
-    //     bytes32[] memory leaves = new bytes32[](2);
-    //     leaves[0] = bytes32(uint256(1));
-    //     leaves[1] = bytes32(uint256(2));
+    function testWriteLeavesToJson() public {
+        bytes32[] memory leaves = new bytes32[](2);
+        leaves[0] = bytes32(uint256(1));
+        leaves[1] = bytes32(uint256(2));
 
-    //     bytes32[] memory tokenLeaves = new bytes32[](2);
-    //     tokenLeaves[0] = bytes32(uint256(3));
-    //     tokenLeaves[1] = bytes32(uint256(4));
+        bytes32[] memory tokenLeaves = new bytes32[](2);
+        tokenLeaves[0] = bytes32(uint256(3));
+        tokenLeaves[1] = bytes32(uint256(4));
 
-    //     SetupPaymentsLib.writeLeavesToJson(leaves, tokenLeaves, vm);
+        SetupPaymentsLib.writeLeavesToJson(leaves, tokenLeaves);
 
-    //     assertTrue(vm.exists("payments.json"), "JSON file should be created");
-    // }
+        assertTrue(vm.exists("payments.json"), "JSON file should be created");
+    }
 
-    // function testParseLeavesFromJson() public {
-    //     string memory filePath = "test_parse_payments.json";
-    //     string memory jsonContent = '{"leaves":["0x1234"], "tokenLeaves":["0x5678"]}';
-    //     vm.writeFile(filePath, jsonContent);
+    function testParseLeavesFromJson() public {
+        string memory filePath = "test_parse_payments.json";
+        string memory jsonContent = '{"leaves":["0x1234"], "tokenLeaves":["0x5678"]}';
+        vm.writeFile(filePath, jsonContent);
 
-    //     SetupPaymentsLib.PaymentLeaves memory paymentLeaves = SetupPaymentsLib.parseLeavesFromJson(filePath, vm);
+        SetupPaymentsLib.PaymentLeaves memory paymentLeaves = SetupPaymentsLib.parseLeavesFromJson(filePath);
 
-    //     assertEq(paymentLeaves.leaves.length, 1, "Incorrect number of leaves");
-    //     assertEq(paymentLeaves.tokenLeaves.length, 1, "Incorrect number of token leaves");
-    // }
+        assertEq(paymentLeaves.leaves.length, 1, "Incorrect number of leaves");
+        assertEq(paymentLeaves.tokenLeaves.length, 1, "Incorrect number of token leaves");
+    }
+
+    function testGenerateMerkleProof() public {
+        SetupPaymentsLib.PaymentLeaves memory paymentLeaves = SetupPaymentsLib.parseLeavesFromJson("test/mockData/scratch/payment_leaves.json");
+
+        bytes32[] memory leaves = paymentLeaves.leaves;
+        uint256 indexToProve = 0;
+
+        bytes32[] memory proof = new bytes32[](2);
+        proof[0] = leaves[1];
+        proof[1] = keccak256(abi.encodePacked(leaves[2], leaves[3]));
+
+        bytes memory proofBytes1 = abi.encodePacked(proof);
+        bytes memory proofBytes2 = SetupPaymentsLib.generateMerkleProof(leaves, indexToProve);
+
+        require(keccak256(proofBytes1) == keccak256(proofBytes2), "Proofs do not match");
+    }
 
      function testProcessClaim() public {
         string memory filePath = "test/mockData/scratch/payment_leaves.json";
@@ -126,9 +142,6 @@ contract SetupPaymentsLibTest is Test, TestConstants {
             amountPerPayment,
             duration
         );
-
-        // The checks are performed inside the MockRewardsCoordinator
-        // If the function doesn't revert, it means all checks passed
     }
 }
 

diff --git a/contracts/test/mockData/scratch/payment_leaves.json b/contracts/test/mockData/scratch/payment_leaves.json
@@ -3,14 +3,12 @@
       "0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef",
       "0x2345678901bcdef2345678901bcdef2345678901bcdef2345678901bcdef2345",
       "0x3456789012cdef3456789012cdef3456789012cdef3456789012cdef3456789",
-      "0x4567890123def4567890123def4567890123def4567890123def4567890123d",
-      "0x5678901234ef5678901234ef5678901234ef5678901234ef5678901234ef567"
+      "0x4567890123def4567890123def4567890123def4567890123def4567890123d"
     ],
     "tokenLeaves": [
       "0xabcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890",
       "0xbcdef2345678901bcdef2345678901bcdef2345678901bcdef2345678901bcde",
       "0xcdef3456789012cdef3456789012cdef3456789012cdef3456789012cdef345",
-      "0xdef4567890123def4567890123def4567890123def4567890123def4567890",
-      "0xef5678901234ef5678901234ef5678901234ef5678901234ef5678901234ef5"
+      "0xdef4567890123def4567890123def4567890123def4567890123def4567890"
     ]
   }
diff --git a/contracts/test_parse_payments.json b/contracts/test_parse_payments.json
@@ -0,0 +1 @@
+{"leaves":["0x1234"], "tokenLeaves":["0x5678"]}