test: improve coverage for GraphPayments

tmigone · tmigone · commit 975e6db876b2 · 2024-12-13T14:23:16.000-03:00
Signed-off-by: Tomás Migone &lt;tomas@edgeandnode.com&gt;
diff --git a/.gitignore b/.gitignore
@@ -34,8 +34,10 @@ bin/
 .vscode
 
 # Coverage and other reports
+coverage/
 reports/
 coverage.json
+lcov.info
 
 # Local test files
 addresses-local.json
diff --git a/packages/horizon/contracts/interfaces/IGraphPayments.sol b/packages/horizon/contracts/interfaces/IGraphPayments.sol
@@ -20,6 +20,7 @@ interface IGraphPayments {
 
     /**
      * @notice Emitted when a payment is collected
+     * @param paymentType The type of payment as defined in {IGraphPayments}
      * @param payer The address of the payer
      * @param receiver The address of the receiver
      * @param dataService The address of the data service
@@ -30,6 +31,7 @@ interface IGraphPayments {
      * @param tokensReceiver Amount of tokens for the receiver
      */
     event GraphPaymentCollected(
+        PaymentTypes paymentType,
         address indexed payer,
         address indexed receiver,
         address indexed dataService,
@@ -40,14 +42,6 @@ interface IGraphPayments {
         uint256 tokensReceiver
     );
 
-    /**
-     * @notice Thrown when the calculated amount of tokens to be paid out to all parties is
-     * not the same as the amount of tokens being collected
-     * @param tokens The amount of tokens being collected
-     * @param tokensCalculated The sum of all the tokens to be paid out
-     */
-    error GraphPaymentsBadAccounting(uint256 tokens, uint256 tokensCalculated);
-
     /**
      * @notice Thrown when the protocol payment cut is invalid
      * @param protocolPaymentCut The protocol payment cut
diff --git a/packages/horizon/contracts/payments/GraphPayments.sol b/packages/horizon/contracts/payments/GraphPayments.sol
@@ -3,6 +3,7 @@ pragma solidity 0.8.27;
 
 import { IGraphToken } from "@graphprotocol/contracts/contracts/token/IGraphToken.sol";
 import { IGraphPayments } from "../interfaces/IGraphPayments.sol";
+import { IHorizonStakingTypes } from "../interfaces/internal/IHorizonStakingTypes.sol";
 
 import { Initializable } from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 import { MulticallUpgradeable } from "@openzeppelin/contracts-upgradeable/utils/MulticallUpgradeable.sol";
@@ -70,14 +71,14 @@ contract GraphPayments is Initializable, MulticallUpgradeable, GraphDirectory, I
         uint256 tokensDataService = tokensRemaining.mulPPMRoundUp(dataServiceCut);
         tokensRemaining = tokensRemaining - tokensDataService;
 
-        uint256 tokensDelegationPool = tokensRemaining.mulPPMRoundUp(
-            _graphStaking().getDelegationFeeCut(receiver, dataService, paymentType)
-        );
-        tokensRemaining = tokensRemaining - tokensDelegationPool;
-
-        // Ensure accounting is correct
-        uint256 tokensTotal = tokensProtocol + tokensDataService + tokensDelegationPool + tokensRemaining;
-        require(tokens == tokensTotal, GraphPaymentsBadAccounting(tokens, tokensTotal));
+        uint256 tokensDelegationPool = 0;
+        IHorizonStakingTypes.DelegationPool memory pool = _graphStaking().getDelegationPool(receiver, dataService);
+        if (pool.shares > 0) {
+            tokensDelegationPool = tokensRemaining.mulPPMRoundUp(
+                _graphStaking().getDelegationFeeCut(receiver, dataService, paymentType)
+            );
+            tokensRemaining = tokensRemaining - tokensDelegationPool;
+        }
 
         // Pay all parties
         _graphToken().burnTokens(tokensProtocol);
@@ -92,6 +93,7 @@ contract GraphPayments is Initializable, MulticallUpgradeable, GraphDirectory, I
         _graphToken().pushTokens(receiver, tokensRemaining);
 
         emit GraphPaymentCollected(
+            paymentType,
             msg.sender,
             receiver,
             dataService,
diff --git a/packages/horizon/test/GraphBase.t.sol b/packages/horizon/test/GraphBase.t.sol
@@ -41,7 +41,7 @@ abstract contract GraphBaseTest is IHorizonStakingTypes, Utils, Constants {
     RewardsManagerMock public rewardsManager;
     CurationMock public curation;
     TAPCollector tapCollector;
-    
+
     HorizonStaking private stakingBase;
     HorizonStakingExtension private stakingExtension;
 
@@ -102,22 +102,34 @@ abstract contract GraphBaseTest is IHorizonStakingTypes, Utils, Constants {
         GraphProxy stakingProxy = new GraphProxy(address(0), address(proxyAdmin));
 
         // GraphPayments predict address
-        bytes memory paymentsParameters = abi.encode(address(controller), protocolPaymentCut);
-        bytes memory paymentsBytecode = abi.encodePacked(
+        bytes memory paymentsImplementationParameters = abi.encode(address(controller), protocolPaymentCut);
+        bytes memory paymentsImplementationBytecode = abi.encodePacked(
             type(GraphPayments).creationCode,
-            paymentsParameters
+            paymentsImplementationParameters
         );
-        address predictedPaymentsAddress = _computeAddress(
+        address predictedPaymentsImplementationAddress = _computeAddress(
             "GraphPayments",
-            paymentsBytecode,
+            paymentsImplementationBytecode,
             users.deployer
         );
-        
-        // PaymentsEscrow
-        bytes memory escrowImplementationParameters = abi.encode(
-            address(controller),
-            withdrawEscrowThawingPeriod
+
+        bytes memory paymentsProxyParameters = abi.encode(
+            predictedPaymentsImplementationAddress,
+            users.governor,
+            abi.encodeCall(GraphPayments.initialize, ())
+        );
+        bytes memory paymentsProxyBytecode = abi.encodePacked(
+            type(TransparentUpgradeableProxy).creationCode,
+            paymentsProxyParameters
+        );
+        address predictedPaymentsProxyAddress = _computeAddress(
+            "TransparentUpgradeableProxy",
+            paymentsProxyBytecode,
+            users.deployer
         );
+
+        // PaymentsEscrow
+        bytes memory escrowImplementationParameters = abi.encode(address(controller), withdrawEscrowThawingPeriod);
         bytes memory escrowImplementationBytecode = abi.encodePacked(
             type(PaymentsEscrow).creationCode,
             escrowImplementationParameters
@@ -156,29 +168,32 @@ abstract contract GraphBaseTest is IHorizonStakingTypes, Utils, Constants {
         resetPrank(users.governor);
         controller.setContractProxy(keccak256("GraphToken"), address(token));
         controller.setContractProxy(keccak256("PaymentsEscrow"), predictedEscrowProxyAddress);
-        controller.setContractProxy(keccak256("GraphPayments"), predictedPaymentsAddress);
+        controller.setContractProxy(keccak256("GraphPayments"), predictedPaymentsProxyAddress);
         controller.setContractProxy(keccak256("Staking"), address(stakingProxy));
         controller.setContractProxy(keccak256("EpochManager"), address(epochManager));
         controller.setContractProxy(keccak256("RewardsManager"), address(rewardsManager));
         controller.setContractProxy(keccak256("Curation"), address(curation));
         controller.setContractProxy(keccak256("GraphTokenGateway"), graphTokenGatewayAddress);
         controller.setContractProxy(keccak256("GraphProxyAdmin"), address(proxyAdmin));
-        
-        resetPrank(users.deployer);
-        address paymentsAddress = _deployContract("GraphPayments", paymentsBytecode);
-        assertEq(paymentsAddress, predictedPaymentsAddress);
-        payments = GraphPayments(paymentsAddress);
-
-        address escrowImplementationAddress = _deployContract("PaymentsEscrow", escrowImplementationBytecode);
-        address escrowProxyAddress = _deployContract("TransparentUpgradeableProxy", escrowProxyBytecode);
-        assertEq(escrowImplementationAddress, predictedEscrowImplementationAddress);
-        assertEq(escrowProxyAddress, predictedEscrowProxyAddress);
-        escrow = PaymentsEscrow(escrowProxyAddress);
 
-        stakingExtension = new HorizonStakingExtension(
-            address(controller),
-            subgraphDataServiceLegacyAddress
-        );
+        resetPrank(users.deployer);
+        {
+            address paymentsImplementationAddress = _deployContract("GraphPayments", paymentsImplementationBytecode);
+            address paymentsProxyAddress = _deployContract("TransparentUpgradeableProxy", paymentsProxyBytecode);
+            assertEq(paymentsImplementationAddress, predictedPaymentsImplementationAddress);
+            assertEq(paymentsProxyAddress, predictedPaymentsProxyAddress);
+            payments = GraphPayments(paymentsProxyAddress);
+        }
+
+        {
+            address escrowImplementationAddress = _deployContract("PaymentsEscrow", escrowImplementationBytecode);
+            address escrowProxyAddress = _deployContract("TransparentUpgradeableProxy", escrowProxyBytecode);
+            assertEq(escrowImplementationAddress, predictedEscrowImplementationAddress);
+            assertEq(escrowProxyAddress, predictedEscrowProxyAddress);
+            escrow = PaymentsEscrow(escrowProxyAddress);
+        }
+
+        stakingExtension = new HorizonStakingExtension(address(controller), subgraphDataServiceLegacyAddress);
         stakingBase = new HorizonStaking(
             address(controller),
             address(stakingExtension),
@@ -228,7 +243,11 @@ abstract contract GraphBaseTest is IHorizonStakingTypes, Utils, Constants {
      * PRIVATE
      */
 
-    function _computeAddress(string memory contractName, bytes memory bytecode, address deployer) private pure returns (address) {
+    function _computeAddress(
+        string memory contractName,
+        bytes memory bytecode,
+        address deployer
+    ) private pure returns (address) {
         bytes32 salt = keccak256(abi.encodePacked(contractName, "Salt"));
         return Create2.computeAddress(salt, keccak256(bytecode), deployer);
     }
@@ -237,4 +256,4 @@ abstract contract GraphBaseTest is IHorizonStakingTypes, Utils, Constants {
         bytes32 salt = keccak256(abi.encodePacked(contractName, "Salt"));
         return Create2.deploy(0, salt, bytecode);
     }
-}
+}
diff --git a/packages/horizon/test/escrow/GraphEscrow.t.sol b/packages/horizon/test/escrow/GraphEscrow.t.sol
@@ -4,6 +4,7 @@ pragma solidity 0.8.27;
 import "forge-std/Test.sol";
 import { IPaymentsEscrow } from "../../contracts/interfaces/IPaymentsEscrow.sol";
 import { IGraphPayments } from "../../contracts/interfaces/IGraphPayments.sol";
+import { IHorizonStakingTypes } from "../../contracts/interfaces/internal/IHorizonStakingTypes.sol";
 
 import { HorizonStakingSharedTest } from "../shared/horizon-staking/HorizonStakingShared.t.sol";
 import { PaymentsEscrowSharedTest } from "../shared/payments-escrow/PaymentsEscrowShared.t.sol";
@@ -112,9 +113,12 @@ contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
         uint256 tokensDataService = (_tokens - _tokens.mulPPMRoundUp(payments.PROTOCOL_PAYMENT_CUT())).mulPPMRoundUp(
             _dataServiceCut
         );
-        uint256 tokensDelegation = (_tokens -
-            _tokens.mulPPMRoundUp(payments.PROTOCOL_PAYMENT_CUT()) -
-            tokensDataService).mulPPMRoundUp(staking.getDelegationFeeCut(_receiver, _dataService, _paymentType));
+        uint256 tokensDelegation = 0;
+        IHorizonStakingTypes.DelegationPool memory pool = staking.getDelegationPool(_receiver, _dataService);
+        if (pool.shares > 0) {
+            tokensDelegation = (_tokens - _tokens.mulPPMRoundUp(payments.PROTOCOL_PAYMENT_CUT()) - tokensDataService)
+                .mulPPMRoundUp(staking.getDelegationFeeCut(_receiver, _dataService, _paymentType));
+        }
         uint256 receiverExpectedPayment = _tokens -
             _tokens.mulPPMRoundUp(payments.PROTOCOL_PAYMENT_CUT()) -
             tokensDataService -
diff --git a/packages/horizon/test/escrow/collect.t.sol b/packages/horizon/test/escrow/collect.t.sol
@@ -15,6 +15,7 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
 
     function testCollect_Tokens(
         uint256 tokens,
+        uint256 tokensToCollect,
         uint256 delegationTokens,
         uint256 dataServiceCut
     )
@@ -25,13 +26,43 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
     {
         dataServiceCut = bound(dataServiceCut, 0, MAX_PPM);
         delegationTokens = bound(delegationTokens, 1, MAX_STAKING_TOKENS);
+        tokensToCollect = bound(tokensToCollect, 1, MAX_STAKING_TOKENS);
 
         resetPrank(users.delegator);
         _delegate(users.indexer, subgraphDataServiceAddress, delegationTokens, 0);
 
+        resetPrank(users.gateway);
+        _depositTokens(users.verifier, users.indexer, tokensToCollect);
+
+        // burn some tokens to prevent overflow
+        resetPrank(users.indexer);
+        token.burn(MAX_STAKING_TOKENS);
+
+        resetPrank(users.verifier);
+        _collectEscrow(
+            IGraphPayments.PaymentTypes.QueryFee,
+            users.gateway,
+            users.indexer,
+            tokensToCollect,
+            subgraphDataServiceAddress,
+            dataServiceCut
+        );
+    }
+
+    function testCollect_Tokens_NoProvision(
+        uint256 tokens,
+        uint256 dataServiceCut
+    ) public useIndexer useDelegationFeeCut(IGraphPayments.PaymentTypes.QueryFee, delegationFeeCut) {
+        dataServiceCut = bound(dataServiceCut, 0, MAX_PPM);
+        tokens = bound(tokens, 1, MAX_STAKING_TOKENS);
+
         resetPrank(users.gateway);
         _depositTokens(users.verifier, users.indexer, tokens);
 
+        // burn some tokens to prevent overflow
+        resetPrank(users.indexer);
+        token.burn(MAX_STAKING_TOKENS);
+
         resetPrank(users.verifier);
         _collectEscrow(
             IGraphPayments.PaymentTypes.QueryFee,
@@ -67,62 +98,4 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
         );
         vm.stopPrank();
     }
-
-    function testCollect_RevertWhen_InvalidPool(
-        uint256 amount
-    )
-        public
-        useIndexer
-        useProvision(amount, 0, 0)
-        useDelegationFeeCut(IGraphPayments.PaymentTypes.QueryFee, delegationFeeCut)
-    {
-        vm.assume(amount > 1 ether);
-
-        resetPrank(users.gateway);
-        _depositTokens(users.verifier, users.indexer, amount);
-
-        resetPrank(users.verifier);
-        vm.expectRevert(
-            abi.encodeWithSelector(
-                IHorizonStakingMain.HorizonStakingInvalidDelegationPool.selector,
-                users.indexer,
-                subgraphDataServiceAddress
-            )
-        );
-        escrow.collect(
-            IGraphPayments.PaymentTypes.QueryFee,
-            users.gateway,
-            users.indexer,
-            amount,
-            subgraphDataServiceAddress,
-            1
-        );
-    }
-
-    function testCollect_RevertWhen_InvalidProvision(
-        uint256 amount
-    ) public useIndexer useDelegationFeeCut(IGraphPayments.PaymentTypes.QueryFee, delegationFeeCut) {
-        vm.assume(amount > 1 ether);
-        vm.assume(amount <= MAX_STAKING_TOKENS);
-
-        resetPrank(users.gateway);
-        _depositTokens(users.verifier, users.indexer, amount);
-
-        resetPrank(users.verifier);
-        vm.expectRevert(
-            abi.encodeWithSelector(
-                IHorizonStakingMain.HorizonStakingInvalidProvision.selector,
-                users.indexer,
-                subgraphDataServiceAddress
-            )
-        );
-        escrow.collect(
-            IGraphPayments.PaymentTypes.QueryFee,
-            users.gateway,
-            users.indexer,
-            amount,
-            subgraphDataServiceAddress,
-            1
-        );
-    }
 }
diff --git a/packages/horizon/test/payments/GraphPayments.t.sol b/packages/horizon/test/payments/GraphPayments.t.sol
