fix: renamed error, move code to trigger revert

graphprotocol · May 23, 2024 · 3b39aa5 · 3b39aa5
1 parent 8947c07
commit 3b39aa5
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 11 deletions.
diff --git a/packages/horizon/contracts/payments/GraphPayments.sol b/packages/horizon/contracts/payments/GraphPayments.sol
@@ -30,7 +30,7 @@ contract GraphPayments is Multicall, GraphDirectory, GraphPaymentsStorageV1Stora
     error GraphPaymentsNotThawing();
     error GraphPaymentsStillThawing(uint256 currentTimestamp, uint256 thawEndTimestamp);
     error GraphPaymentsCollectorNotAuthorized(address sender, address dataService);
-    error GraphPaymentsCollectorInsufficientAmount(uint256 available, uint256 required);
+    error GraphPaymentsInsufficientAmount(uint256 available, uint256 required);
 
     // -- Events --
 
@@ -54,25 +54,27 @@ contract GraphPayments is Multicall, GraphDirectory, GraphPaymentsStorageV1Stora
     ) external {
         _graphToken().pullTokens(msg.sender, tokens);
 
-        // Pay protocol cut
+        // Calculate cuts
         uint256 tokensProtocol = (tokens * PROTOCOL_PAYMENT_CUT) / MAX_PPM;
+        uint256 delegationFeeCut = _graphStaking().getDelegationFeeCut(receiver, dataService, uint8(paymentType));
+        uint256 tokensDelegationPool = (tokens * delegationFeeCut) / MAX_PPM;
+        uint256 totalCut = tokensProtocol + tokensDataService + tokensDelegationPool;
+        if (totalCut > tokens) {
+            revert GraphPaymentsInsufficientAmount(tokens, totalCut);
+        }
+
+        // Pay protocol cut
         _graphToken().burnTokens(tokensProtocol);
 
         // Pay data service cut
         _graphToken().pushTokens(dataService, tokensDataService);
 
-        // Get delegation cut
-        uint256 delegationFeeCut = _graphStaking().getDelegationFeeCut(receiver, dataService, uint8(paymentType));
-        uint256 tokensDelegationPool = (tokens * delegationFeeCut) / MAX_PPM;
+        // Pay delegators
         if (tokensDelegationPool > 0) {
             _graphStaking().addToDelegationPool(receiver, dataService, tokensDelegationPool);
         }
 
-        // Pay the rest to the receiver
-        uint256 totalCut = tokensProtocol + tokensDataService + tokensDelegationPool;
-        if (totalCut > tokens) {
-            revert GraphPaymentsCollectorInsufficientAmount(tokens, totalCut);
-        }
+        // Pay receiver
         uint256 tokensReceiverRemaining = tokens - totalCut;
         _graphToken().pushTokens(receiver, tokensReceiverRemaining);
 

diff --git a/packages/horizon/test/payments/GraphPayments.t.sol b/packages/horizon/test/payments/GraphPayments.t.sol
@@ -9,7 +9,10 @@ import { HorizonStakingSharedTest } from "../shared/horizon-staking/HorizonStaki
 
 contract GraphPaymentsTest is HorizonStakingSharedTest {
 
-    function testCollect(uint256 amount, uint256 tokensDataService) public useProvision(amount, 0, 0) useDelegationFeeCut(0, 100000) {
+    function testCollect(
+        uint256 amount,
+        uint256 tokensDataService
+    ) public useProvision(amount, 0, 0) useDelegationFeeCut(0, delegationFeeCut) {
         uint256 tokensProtocol = amount * protocolPaymentCut / MAX_PPM;
         uint256 tokensDelegatoion = amount * delegationFeeCut / MAX_PPM;
         vm.assume(tokensDataService < amount - tokensProtocol - tokensDelegatoion);
@@ -35,4 +38,27 @@ contract GraphPaymentsTest is HorizonStakingSharedTest {
         uint256 delegatorBalance = staking.getDelegatedTokensAvailable(users.indexer, subgraphDataServiceAddress);
         assertEq(delegatorBalance, tokensDelegatoion);
     }
+
+    function testCollect_RevertWhen_InsufficientAmount(
+        uint256 amount,
+        uint256 tokensDataService
+    ) public useProvision(amount, 0, 0) useDelegationFeeCut(0, delegationFeeCut) {
+        vm.assume(tokensDataService <= 10_000_000_000 ether);
+        vm.assume(tokensDataService > amount);
+
+        address escrowAddress = address(escrow);
+        mint(escrowAddress, amount);
+        vm.startPrank(escrowAddress);
+        approve(address(payments), amount);
+
+        bytes memory expectedError;
+        {
+            uint256 tokensProtocol = amount * protocolPaymentCut / MAX_PPM;
+            uint256 tokensDelegatoion = amount * delegationFeeCut / MAX_PPM;
+            uint256 requiredAmount = tokensDataService + tokensProtocol + tokensDelegatoion;
+            expectedError = abi.encodeWithSignature("GraphPaymentsInsufficientAmount(uint256,uint256)", amount, requiredAmount);
+        }
+        vm.expectRevert(expectedError);
+        payments.collect(IGraphPayments.PaymentTypes.QueryFee, users.indexer, amount, subgraphDataServiceAddress, tokensDataService);
+    }
 }