[WIP] Graph Horizon: Staking unit tests

packages/horizon/contracts/staking/HorizonStaking.sol

@@ -525,7 +525,7 @@ contract HorizonStaking is HorizonStakingBase, IHorizonStakingMain {
     }
 
     function setMaxThawingPeriod(uint64 maxThawingPeriod) external override onlyGovernor {
-        maxThawingPeriod = _maxThawingPeriod;
+        _maxThawingPeriod = maxThawingPeriod;
         emit MaxThawingPeriodSet(_maxThawingPeriod);
     }
 

packages/horizon/test/GraphBase.t.sol

@@ -54,11 +54,13 @@ abstract contract GraphBaseTest is Test, Constants {
             indexer: createUser("indexer"),
             operator: createUser("operator"),
             gateway: createUser("gateway"),
-            verifier: createUser("verifier")
+            verifier: createUser("verifier"),
+            delegator: createUser("delegator")
         });
 
         // Deploy protocol contracts
         deployProtocolContracts();
+        setupProtocol();
         unpauseProtocol();
 
         // Label contracts
@@ -147,11 +149,15 @@ abstract contract GraphBaseTest is Test, Constants {
         proxyAdmin.upgrade(stakingProxy, address(stakingBase));
         proxyAdmin.acceptProxy(stakingBase, stakingProxy);
         staking = IHorizonStaking(address(stakingProxy));
-        vm.stopPrank();
+    }
+
+    function setupProtocol() private {
+        vm.startPrank(users.governor);
+        staking.setMaxThawingPeriod(MAX_THAWING_PERIOD);
     }
 
     function unpauseProtocol() private {
-        vm.prank(users.governor);
+        vm.startPrank(users.governor);
         controller.setPaused(false);
     }
 

packages/horizon/test/escrow/GraphEscrow.t.sol

@@ -3,7 +3,7 @@ pragma solidity 0.8.26;
 
 import "forge-std/Test.sol";
 
-import { HorizonStakingSharedTest } from "../shared/horizon-staking/HorizonStaking.t.sol";
+import { HorizonStakingSharedTest } from "../shared/horizon-staking/HorizonStakingShared.t.sol";
 
 contract GraphEscrowTest is HorizonStakingSharedTest {
 
@@ -14,21 +14,19 @@ contract GraphEscrowTest is HorizonStakingSharedTest {
     }
 
     modifier approveEscrow(uint256 tokens) {
-        changePrank(users.gateway);
         _approveEscrow(tokens);
         _;
     }
 
     modifier useDeposit(uint256 tokens) {
-        changePrank(users.gateway);
         vm.assume(tokens > 0);
         vm.assume(tokens <= 10_000_000_000 ether);
         _depositTokens(tokens);
         _;
     }
 
     modifier useCollector(uint256 tokens) {
-        changePrank(users.gateway);
+        vm.assume(tokens > 0);
         escrow.approveCollector(users.verifier, tokens);
         _;
     }

packages/horizon/test/escrow/collect.t.sol

@@ -8,18 +8,20 @@ import { IGraphPayments } from "../../contracts/interfaces/IGraphPayments.sol";
 
 contract GraphEscrowCollectTest is GraphEscrowTest {
 
-    function testCollect_Tokens(uint256 amount, uint256 tokensDataService) public useProvision(amount, 0, 0) useDelegationFeeCut(IGraphPayments.PaymentTypes.QueryFee, delegationFeeCut) {
+    function testCollect_Tokens(
+        uint256 amount,
+        uint256 tokensDataService
+    ) public useIndexer useProvision(amount, 0, 0) useDelegationFeeCut(IGraphPayments.PaymentTypes.QueryFee, delegationFeeCut) {
         uint256 tokensProtocol = amount * protocolPaymentCut / MAX_PPM;
         uint256 tokensDelegatoion = amount * delegationFeeCut / MAX_PPM;
         vm.assume(tokensDataService < amount - tokensProtocol - tokensDelegatoion);
 
         vm.startPrank(users.gateway);
         escrow.approveCollector(users.verifier, amount);
         _depositTokens(amount);
-        vm.stopPrank();
 
         uint256 indexerPreviousBalance = token.balanceOf(users.indexer);
-        vm.prank(users.verifier);
+        vm.startPrank(users.verifier);
         escrow.collect(IGraphPayments.PaymentTypes.QueryFee, users.gateway, users.indexer, amount, subgraphDataServiceAddress, tokensDataService);
 
         uint256 indexerBalance = token.balanceOf(users.indexer);
@@ -31,7 +33,7 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
     function testCollect_RevertWhen_CollectorNotAuthorized(uint256 amount) public {
         vm.startPrank(users.verifier);
         uint256 dataServiceCut = 30000; // 3%
-        bytes memory expectedError = abi.encodeWithSignature("GraphEscrowCollectorNotAuthorized(address,address)", users.gateway, users.verifier);
+        bytes memory expectedError = abi.encodeWithSignature("PaymentsEscrowCollectorNotAuthorized(address,address)", users.gateway, users.verifier);
         vm.expectRevert(expectedError);
         escrow.collect(IGraphPayments.PaymentTypes.QueryFee, users.gateway, users.indexer, amount, subgraphDataServiceAddress, dataServiceCut);
         vm.stopPrank();
@@ -43,8 +45,12 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
     ) public useGateway useCollector(insufficientAmount) useDeposit(amount) {
         vm.assume(insufficientAmount < amount);
 
-        changePrank(users.verifier);
-        bytes memory expectedError = abi.encodeWithSignature("GraphEscrowInsufficientAllowance(uint256,uint256)", insufficientAmount, amount);
+        vm.startPrank(users.verifier);
+        bytes memory expectedError = abi.encodeWithSignature(
+            "PaymentsEscrowInsufficientAllowance(uint256,uint256)", 
+            insufficientAmount, 
+            amount
+        );
         vm.expectRevert(expectedError);
         escrow.collect(IGraphPayments.PaymentTypes.QueryFee, users.gateway, users.indexer, amount, subgraphDataServiceAddress, 0);
     }
@@ -55,8 +61,8 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
     ) public useGateway useCollector(amount) useDeposit(insufficientAmount)  {
         vm.assume(insufficientAmount < amount);
 
-        changePrank(users.verifier);
-        bytes memory expectedError = abi.encodeWithSignature("GraphEscrowInsufficientBalance(uint256,uint256)", insufficientAmount, amount);
+        vm.startPrank(users.verifier);
+        bytes memory expectedError = abi.encodeWithSignature("PaymentsEscrowInsufficientBalance(uint256,uint256)", insufficientAmount, amount);
         vm.expectRevert(expectedError);
         escrow.collect(IGraphPayments.PaymentTypes.QueryFee, users.gateway, users.indexer, amount, subgraphDataServiceAddress, 0);
         vm.stopPrank();

packages/horizon/test/escrow/collector.t.sol

@@ -20,7 +20,7 @@ contract GraphEscrowCollectorTest is GraphEscrowTest {
         uint256 smallerAmount
     ) public useGateway useCollector(amount) {
         vm.assume(smallerAmount < amount);
-        bytes memory expectedError = abi.encodeWithSignature("GraphEscrowInsufficientAllowance(uint256,uint256)", amount, smallerAmount);
+        bytes memory expectedError = abi.encodeWithSignature("PaymentsEscrowInconsistentAllowance(uint256,uint256)", amount, smallerAmount);
         vm.expectRevert(expectedError);
         escrow.approveCollector(users.verifier, smallerAmount);
     }
@@ -52,7 +52,7 @@ contract GraphEscrowCollectorTest is GraphEscrowTest {
     }
 
     function testCollector_RevertWhen_CancelThawIsNotThawing(uint256 amount) public useGateway useCollector(amount) {
-        bytes memory expectedError = abi.encodeWithSignature("GraphEscrowNotThawing()");
+        bytes memory expectedError = abi.encodeWithSignature("PaymentsEscrowNotThawing()");
         vm.expectRevert(expectedError);
         escrow.cancelThawCollector(users.verifier);
         vm.stopPrank();
@@ -70,14 +70,14 @@ contract GraphEscrowCollectorTest is GraphEscrowTest {
     }
 
     function testCollector_RevertWhen_RevokeIsNotThawing(uint256 amount) public useGateway useCollector(amount) {
-        bytes memory expectedError = abi.encodeWithSignature("GraphEscrowNotThawing()");
+        bytes memory expectedError = abi.encodeWithSignature("PaymentsEscrowNotThawing()");
         vm.expectRevert(expectedError);
         escrow.revokeCollector(users.verifier);
     }
 
     function testCollector_RevertWhen_RevokeIsStillThawing(uint256 amount) public useGateway useCollector(amount) {
         escrow.thawCollector(users.verifier);
-        bytes memory expectedError = abi.encodeWithSignature("GraphEscrowStillThawing(uint256,uint256)", block.timestamp, block.timestamp + revokeCollectorThawingPeriod);
+        bytes memory expectedError = abi.encodeWithSignature("PaymentsEscrowStillThawing(uint256,uint256)", block.timestamp, block.timestamp + revokeCollectorThawingPeriod);
         vm.expectRevert(expectedError);
         escrow.revokeCollector(users.verifier);
     }

packages/horizon/test/escrow/thaw.t.sol

@@ -18,7 +18,7 @@ contract GraphEscrowThawTest is GraphEscrowTest {
     function testThaw_RevertWhen_InsufficientThawAmount(
         uint256 amount
     ) public useGateway useDeposit(amount) {
-        bytes memory expectedError = abi.encodeWithSignature("GraphEscrowInsufficientTokensThawing()");
+        bytes memory expectedError = abi.encodeWithSignature("PaymentsEscrowInsufficientTokensThawing()");
         vm.expectRevert(expectedError);
         escrow.thaw(users.indexer, 0);
     }
@@ -27,7 +27,7 @@ contract GraphEscrowThawTest is GraphEscrowTest {
         uint256 amount
     ) public useGateway useDeposit(amount) {
         uint256 overAmount = amount + 1;
-        bytes memory expectedError = abi.encodeWithSignature("GraphEscrowInsufficientBalance(uint256,uint256)", amount, overAmount);
+        bytes memory expectedError = abi.encodeWithSignature("PaymentsEscrowInsufficientBalance(uint256,uint256)", amount, overAmount);
         vm.expectRevert(expectedError);
         escrow.thaw(users.indexer, overAmount);
     }

packages/horizon/test/escrow/withdraw.t.sol

@@ -30,19 +30,17 @@ contract GraphEscrowWithdrawTest is GraphEscrowTest {
     }
 
     function testWithdraw_RevertWhen_NotThawing(uint256 amount) public useGateway useDeposit(amount) {
-        bytes memory expectedError = abi.encodeWithSignature("GraphEscrowNotThawing()");
+        bytes memory expectedError = abi.encodeWithSignature("PaymentsEscrowNotThawing()");
         vm.expectRevert(expectedError);
         escrow.withdraw(users.indexer);
-        vm.stopPrank();
     }
 
     function testWithdraw_RevertWhen_StillThawing(
         uint256 amount,
         uint256 thawAmount
     ) public useGateway depositAndThawTokens(amount, thawAmount) {
-        bytes memory expectedError = abi.encodeWithSignature("GraphEscrowStillThawing(uint256,uint256)", block.timestamp, block.timestamp + withdrawEscrowThawingPeriod);
+        bytes memory expectedError = abi.encodeWithSignature("PaymentsEscrowStillThawing(uint256,uint256)", block.timestamp, block.timestamp + withdrawEscrowThawingPeriod);
         vm.expectRevert(expectedError);
         escrow.withdraw(users.indexer);
-        vm.stopPrank();
     }
 }

packages/horizon/test/payments/GraphPayments.t.sol

@@ -5,14 +5,14 @@ import "forge-std/Test.sol";
 
 import { IGraphPayments } from "../../contracts/interfaces/IGraphPayments.sol";
 
-import { HorizonStakingSharedTest } from "../shared/horizon-staking/HorizonStaking.t.sol";
+import { HorizonStakingSharedTest } from "../shared/horizon-staking/HorizonStakingShared.t.sol";
 
 contract GraphPaymentsTest is HorizonStakingSharedTest {
 
     function testCollect(
         uint256 amount,
         uint256 tokensDataService
-    ) public useProvision(amount, 0, 0) useDelegationFeeCut(IGraphPayments.PaymentTypes.QueryFee, delegationFeeCut) {
+    ) public useIndexer useProvision(amount, 0, 0) useDelegationFeeCut(IGraphPayments.PaymentTypes.QueryFee, delegationFeeCut) {
         uint256 tokensProtocol = amount * protocolPaymentCut / MAX_PPM;
         uint256 tokensDelegatoion = amount * delegationFeeCut / MAX_PPM;
         vm.assume(tokensDataService < amount - tokensProtocol - tokensDelegatoion);
@@ -42,7 +42,7 @@ contract GraphPaymentsTest is HorizonStakingSharedTest {
     function testCollect_RevertWhen_InsufficientAmount(
         uint256 amount,
         uint256 tokensDataService
-    ) public useProvision(amount, 0, 0) useDelegationFeeCut(IGraphPayments.PaymentTypes.QueryFee, delegationFeeCut) {
+    ) public useIndexer useProvision(amount, 0, 0) useDelegationFeeCut(IGraphPayments.PaymentTypes.QueryFee, delegationFeeCut) {
         vm.assume(tokensDataService <= 10_000_000_000 ether);
         vm.assume(tokensDataService > amount);
 

packages/horizon/test/shared/horizon-staking/HorizonStaking.t.sol → packages/horizon/test/shared/horizon-staking/HorizonStakingShared.t.sol

@@ -8,10 +8,24 @@ import { IGraphPayments } from "../../../contracts/interfaces/IGraphPayments.sol
 
 abstract contract HorizonStakingSharedTest is GraphBaseTest {
 
+    modifier useIndexer() {
+        vm.startPrank(users.indexer);
+        _;
+        vm.stopPrank();
+    }
+
+    modifier assumeProvisionTokens(uint256 tokens) {
+        vm.assume(tokens > MIN_PROVISION_SIZE);
+        vm.assume(tokens <= 10_000_000_000 ether);
+        _;
+    }
+
     modifier useProvision(uint256 tokens, uint32 maxVerifierCut, uint64 thawingPeriod) {
         vm.assume(tokens <= 10_000_000_000 ether);
-        vm.assume(tokens > 1e18);
-        _createProvision(tokens, maxVerifierCut, thawingPeriod);
+        vm.assume(tokens > MIN_PROVISION_SIZE);
+        vm.assume(maxVerifierCut <= MAX_MAX_VERIFIER_CUT);
+        vm.assume(thawingPeriod <= MAX_THAWING_PERIOD);
+        _createProvision(subgraphDataServiceAddress, tokens, maxVerifierCut, thawingPeriod);
         _;
     }
 
@@ -28,22 +42,24 @@ abstract contract HorizonStakingSharedTest is GraphBaseTest {
 
     /* Helpers */
 
-    function _createProvision(uint256 tokens, uint32 maxVerifierCut, uint64 thawingPeriod) internal {
-        vm.startPrank(users.indexer);
+    function _createProvision(
+        address dataServiceAddress,
+        uint256 tokens,
+        uint32 maxVerifierCut,
+        uint64 thawingPeriod
+    ) internal {
         token.approve(address(staking), tokens);
         staking.stakeTo(users.indexer, tokens);
         staking.provision(
             users.indexer,
-            subgraphDataServiceAddress,
+            dataServiceAddress,
             tokens,
             maxVerifierCut,
             thawingPeriod
         );
-        vm.stopPrank();
     }
 
     function _setDelegationFeeCut(IGraphPayments.PaymentTypes paymentType, uint256 cut) internal {
-        vm.prank(users.indexer);
         staking.setDelegationFeeCut(users.indexer, subgraphDataServiceAddress, paymentType, cut);
     }
 }

packages/horizon/test/staking/HorizonStaking.t.sol

@@ -0,0 +1,86 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.26;
+
+import "forge-std/Test.sol";
+
+import { IHorizonStakingTypes } from "../../contracts/interfaces/internal/IHorizonStakingTypes.sol";
+import { HorizonStakingSharedTest } from "../shared/horizon-staking/HorizonStakingShared.t.sol";
+
+contract HorizonStakingTest is HorizonStakingSharedTest, IHorizonStakingTypes {
+
+    modifier useOperator() {
+        vm.startPrank(users.indexer);
+        staking.setOperator(users.operator, subgraphDataServiceAddress, true);
+        vm.startPrank(users.operator);
+        _;
+        vm.stopPrank();
+    }
+
+    modifier useStake(uint256 amount) {
+        vm.assume(amount > MIN_PROVISION_SIZE);
+        approve(address(staking), amount);
+        staking.stake(amount);
+        _;
+    }
+
+    modifier useStakeTo(address to, uint256 amount) {
+        vm.assume(amount > MIN_PROVISION_SIZE);
+        _stakeTo(to, amount);
+        _;
+    }
+
+    modifier useThawRequest(uint256 thawAmount) {
+        vm.assume(thawAmount > 0);
+        _createThawRequest(thawAmount);
+        _;
+    }
+
+    modifier useThawAndDeprovision(uint256 amount, uint64 thawingPeriod) {
+        vm.assume(amount > 0);
+        _createThawRequest(amount);
+        skip(thawingPeriod + 1);
+        _deprovision(amount);
+        _;
+    }
+
+    modifier useDelegation(uint256 delegationAmount) {
+        address msgSender;
+        (, msgSender,) = vm.readCallers();
+        vm.assume(delegationAmount > MIN_DELEGATION);
+        vm.assume(delegationAmount <= 10_000_000_000 ether);
+        vm.startPrank(users.delegator);
+        _delegate(delegationAmount);
+        vm.startPrank(msgSender);
+        _;
+    }
+
+    function _stakeTo(address to, uint256 amount) internal {
+        approve(address(staking), amount);
+        staking.stakeTo(to, amount);
+    }
+
+    function _createThawRequest(uint256 thawAmount) internal returns (bytes32) {
+        return staking.thaw(users.indexer, subgraphDataServiceAddress, thawAmount);
+    }
+
+    function _deprovision(uint256 amount) internal {
+        staking.deprovision(users.indexer, subgraphDataServiceAddress, amount);
+    }
+
+    function _delegate(uint256 amount) internal {
+        token.approve(address(staking), amount);
+        staking.delegate(users.indexer, subgraphDataServiceAddress, amount, 0);
+    }
+
+    function _getDelegation() internal view returns (Delegation memory) {
+        return staking.getDelegation(users.delegator, users.indexer, subgraphDataServiceAddress);
+    }
+
+    function _undelegate(uint256 shares) internal {
+        staking.undelegate(users.indexer, subgraphDataServiceAddress, shares);
+    }
+
+    function _getDelegationPool() internal view returns (DelegationPool memory) {
+        return staking.getDelegationPool(users.indexer, subgraphDataServiceAddress);
+    }
+}