test: improve horizon and subgraph-service test coverage

.gitignore

@@ -105,6 +105,9 @@ tx-builder-*.json
 **/subgraph-service-localNetwork/
 !**/ignition/**/artifacts/
 
+# Temporary test working directories
+**/testing-coverage/
+
 # Claude AI settings
 .claude/
 

packages/horizon/contracts/payments/PaymentsEscrow.sol

@@ -140,6 +140,14 @@ contract PaymentsEscrow is Initializable, MulticallUpgradeable, GraphDirectory,
         // Reduce amount from account balance
         account.balance -= tokens;
 
+        // Cap tokensThawing to the new balance to keep state consistent
+        if (account.tokensThawing > account.balance) {
+            account.tokensThawing = account.balance;
+            if (account.tokensThawing == 0) {
+                account.thawEndTimestamp = 0;
+            }
+        }
+
         uint256 escrowBalanceBefore = _graphToken().balanceOf(address(this));
 
         _graphToken().approve(address(_graphPayments()), tokens);

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

@@ -2,6 +2,7 @@
 pragma solidity 0.8.27;
 
 import { IGraphPayments } from "@graphprotocol/interfaces/contracts/horizon/IGraphPayments.sol";
+import { IPaymentsEscrow } from "@graphprotocol/interfaces/contracts/horizon/IPaymentsEscrow.sol";
 
 import { GraphEscrowTest } from "./GraphEscrow.t.sol";
 
@@ -130,4 +131,134 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
             users.indexer
         );
     }
+
+    function testCollect_EntireBalance(uint256 tokens) public useIndexer {
+        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,
+            users.gateway,
+            users.indexer,
+            tokens,
+            subgraphDataServiceAddress,
+            0,
+            users.indexer
+        );
+
+        // Balance should be zero
+        (uint256 balance, , ) = escrow.escrowAccounts(users.gateway, users.verifier, users.indexer);
+        assertEq(balance, 0);
+    }
+
+    function testCollect_CapsTokensThawingToZero_ResetsThawEndTimestamp(uint256 tokens) public useIndexer {
+        // When collecting the entire balance while thawing, tokensThawing should cap to 0
+        // and thawEndTimestamp should reset to 0
+        tokens = bound(tokens, 1, MAX_STAKING_TOKENS);
+
+        resetPrank(users.gateway);
+        _depositTokens(users.verifier, users.indexer, tokens);
+        escrow.thaw(users.verifier, users.indexer, tokens);
+
+        // burn some tokens to prevent overflow
+        resetPrank(users.indexer);
+        token.burn(MAX_STAKING_TOKENS);
+
+        // Collect entire balance
+        resetPrank(users.verifier);
+        _collectEscrow(
+            IGraphPayments.PaymentTypes.QueryFee,
+            users.gateway,
+            users.indexer,
+            tokens,
+            subgraphDataServiceAddress,
+            0,
+            users.indexer
+        );
+
+        // tokensThawing and thawEndTimestamp should be reset
+        (uint256 balance, uint256 tokensThawingResult, uint256 thawEndTimestamp) = escrow.escrowAccounts(
+            users.gateway,
+            users.verifier,
+            users.indexer
+        );
+        assertEq(balance, 0);
+        assertEq(tokensThawingResult, 0, "tokensThawing should be capped to 0");
+        assertEq(thawEndTimestamp, 0, "thawEndTimestamp should reset when tokensThawing is 0");
+    }
+
+    function testCollect_CapsTokensThawingBelowBalance(uint256 depositAmount, uint256 collectAmount) public useIndexer {
+        // When collecting reduces balance below tokensThawing, tokensThawing should cap at balance
+        depositAmount = bound(depositAmount, 3, MAX_STAKING_TOKENS);
+        collectAmount = bound(collectAmount, 1, depositAmount - 1);
+
+        resetPrank(users.gateway);
+        _depositTokens(users.verifier, users.indexer, depositAmount);
+        // Thaw entire balance
+        escrow.thaw(users.verifier, users.indexer, depositAmount);
+
+        // burn some tokens to prevent overflow
+        resetPrank(users.indexer);
+        token.burn(MAX_STAKING_TOKENS);
+
+        // Collect partial amount
+        resetPrank(users.verifier);
+        _collectEscrow(
+            IGraphPayments.PaymentTypes.QueryFee,
+            users.gateway,
+            users.indexer,
+            collectAmount,
+            subgraphDataServiceAddress,
+            0,
+            users.indexer
+        );
+
+        (uint256 balance, uint256 tokensThawingResult, ) = escrow.escrowAccounts(
+            users.gateway,
+            users.verifier,
+            users.indexer
+        );
+        uint256 remainingBalance = depositAmount - collectAmount;
+        assertEq(balance, remainingBalance);
+        assertEq(tokensThawingResult, remainingBalance, "tokensThawing should cap at remaining balance");
+    }
+
+    function testCollect_RevertWhen_InconsistentCollection(uint256 tokens) public useGateway {
+        tokens = bound(tokens, 1, MAX_STAKING_TOKENS);
+        _depositTokens(users.verifier, users.indexer, tokens);
+
+        // Mock GraphPayments.collect to be a no-op: it succeeds but doesn't pull tokens,
+        // causing the escrow balance to remain unchanged and triggering the consistency check.
+        vm.mockCall(address(payments), abi.encodeWithSelector(IGraphPayments.collect.selector), abi.encode());
+
+        uint256 escrowBalance = token.balanceOf(address(escrow));
+
+        resetPrank(users.verifier);
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                IPaymentsEscrow.PaymentsEscrowInconsistentCollection.selector,
+                escrowBalance,
+                escrowBalance, // balance unchanged because mock didn't pull tokens
+                tokens
+            )
+        );
+        escrow.collect(
+            IGraphPayments.PaymentTypes.QueryFee,
+            users.gateway,
+            users.indexer,
+            tokens,
+            subgraphDataServiceAddress,
+            0,
+            users.indexer
+        );
+
+        vm.clearMockedCalls();
+    }
 }

packages/horizon/test/unit/escrow/constructor.t.sol

@@ -0,0 +1,49 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.27;
+
+import { Test } from "forge-std/Test.sol";
+import { Controller } from "@graphprotocol/contracts/contracts/governance/Controller.sol";
+import { PaymentsEscrow } from "contracts/payments/PaymentsEscrow.sol";
+import { IPaymentsEscrow } from "@graphprotocol/interfaces/contracts/horizon/IPaymentsEscrow.sol";
+
+contract GraphEscrowConstructorTest is Test {
+    Controller public controller;
+
+    function setUp() public {
+        controller = new Controller();
+
+        // GraphDirectory requires all proxy contracts to be registered
+        controller.setContractProxy(keccak256("GraphToken"), makeAddr("GraphToken"));
+        controller.setContractProxy(keccak256("Staking"), makeAddr("Staking"));
+        controller.setContractProxy(keccak256("GraphPayments"), makeAddr("GraphPayments"));
+        controller.setContractProxy(keccak256("PaymentsEscrow"), makeAddr("PaymentsEscrow"));
+        controller.setContractProxy(keccak256("EpochManager"), makeAddr("EpochManager"));
+        controller.setContractProxy(keccak256("RewardsManager"), makeAddr("RewardsManager"));
+        controller.setContractProxy(keccak256("GraphTokenGateway"), makeAddr("GraphTokenGateway"));
+        controller.setContractProxy(keccak256("GraphProxyAdmin"), makeAddr("GraphProxyAdmin"));
+        controller.setContractProxy(keccak256("Curation"), makeAddr("Curation"));
+    }
+
+    function testConstructor_MaxWaitPeriodBoundary() public {
+        uint256 maxWaitPeriod = 90 days;
+
+        // Exactly at MAX_WAIT_PERIOD should succeed
+        PaymentsEscrow escrowAtMax = new PaymentsEscrow(address(controller), maxWaitPeriod);
+        assertEq(escrowAtMax.WITHDRAW_ESCROW_THAWING_PERIOD(), maxWaitPeriod);
+    }
+
+    function testConstructor_RevertWhen_ThawingPeriodTooLong() public {
+        uint256 maxWaitPeriod = 90 days;
+        uint256 tooLong = maxWaitPeriod + 1;
+
+        vm.expectRevert(
+            abi.encodeWithSelector(IPaymentsEscrow.PaymentsEscrowThawingPeriodTooLong.selector, tooLong, maxWaitPeriod)
+        );
+        new PaymentsEscrow(address(controller), tooLong);
+    }
+
+    function testConstructor_ZeroThawingPeriod() public {
+        PaymentsEscrow escrowZero = new PaymentsEscrow(address(controller), 0);
+        assertEq(escrowZero.WITHDRAW_ESCROW_THAWING_PERIOD(), 0);
+    }
+}

packages/horizon/test/unit/escrow/isolation.t.sol

@@ -0,0 +1,76 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.27;
+
+import { GraphEscrowTest } from "./GraphEscrow.t.sol";
+
+contract GraphEscrowIsolationTest is GraphEscrowTest {
+    /*
+     * TESTS
+     */
+
+    function testIsolation_DifferentCollectorsSamePayerReceiver(uint256 amount) public useGateway {
+        amount = bound(amount, 1, MAX_STAKING_TOKENS / 2);
+
+        address collector1 = users.verifier;
+        address collector2 = users.operator;
+
+        _depositTokens(collector1, users.indexer, amount);
+        _depositTokens(collector2, users.indexer, amount * 2);
+
+        (uint256 balance1, , ) = escrow.escrowAccounts(users.gateway, collector1, users.indexer);
+        (uint256 balance2, , ) = escrow.escrowAccounts(users.gateway, collector2, users.indexer);
+
+        assertEq(balance1, amount);
+        assertEq(balance2, amount * 2);
+    }
+
+    function testIsolation_DifferentReceiversSamePayerCollector(uint256 amount) public useGateway {
+        amount = bound(amount, 1, MAX_STAKING_TOKENS / 2);
+
+        address receiver1 = users.indexer;
+        address receiver2 = users.delegator;
+
+        _depositTokens(users.verifier, receiver1, amount);
+        _depositTokens(users.verifier, receiver2, amount * 2);
+
+        (uint256 balance1, , ) = escrow.escrowAccounts(users.gateway, users.verifier, receiver1);
+        (uint256 balance2, , ) = escrow.escrowAccounts(users.gateway, users.verifier, receiver2);
+
+        assertEq(balance1, amount);
+        assertEq(balance2, amount * 2);
+    }
+
+    function testIsolation_ThawOneTupleDoesNotAffectAnother(uint256 amount) public useGateway {
+        amount = bound(amount, 2, MAX_STAKING_TOKENS / 2);
+
+        _depositTokens(users.verifier, users.indexer, amount);
+        _depositTokens(users.verifier, users.delegator, amount);
+
+        // Thaw only the first tuple
+        escrow.thaw(users.verifier, users.indexer, amount / 2);
+
+        // Second tuple should be unaffected
+        (, uint256 tokensThawing2, uint256 thawEndTimestamp2) = escrow.escrowAccounts(
+            users.gateway,
+            users.verifier,
+            users.delegator
+        );
+        assertEq(tokensThawing2, 0);
+        assertEq(thawEndTimestamp2, 0);
+
+        // First tuple should have thawing
+        (, uint256 tokensThawing1, ) = escrow.escrowAccounts(users.gateway, users.verifier, users.indexer);
+        assertEq(tokensThawing1, amount / 2);
+    }
+
+    function testIsolation_EscrowAccounts_NeverUsedAccount() public view {
+        (uint256 balance, uint256 tokensThawing, uint256 thawEndTimestamp) = escrow.escrowAccounts(
+            address(0xdead),
+            address(0xbeef),
+            address(0xface)
+        );
+        assertEq(balance, 0);
+        assertEq(tokensThawing, 0);
+        assertEq(thawEndTimestamp, 0);
+    }
+}

packages/horizon/test/unit/escrow/paused.t.sol

@@ -42,6 +42,14 @@ contract GraphEscrowPausedTest is GraphEscrowTest {
         escrow.thaw(users.verifier, users.indexer, tokens);
     }
 
+    function testPaused_RevertWhen_CancelThaw(
+        uint256 tokens,
+        uint256 thawAmount
+    ) public useGateway depositAndThawTokens(tokens, thawAmount) usePaused(true) {
+        vm.expectRevert(abi.encodeWithSelector(IPaymentsEscrow.PaymentsEscrowIsPaused.selector));
+        escrow.cancelThaw(users.verifier, users.indexer);
+    }
+
     function testPaused_RevertWhen_WithdrawTokens(
         uint256 tokens,
         uint256 thawAmount

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

@@ -39,13 +39,23 @@ contract GraphEscrowWithdrawTest is GraphEscrowTest {
         escrow.withdraw(users.verifier, users.indexer);
     }
 
+    function testWithdraw_SucceedsOneSecondAfterThawEnd(
+        uint256 amount,
+        uint256 thawAmount
+    ) public useGateway depositAndThawTokens(amount, thawAmount) {
+        // Advance time to exactly one second past thaw end
+        skip(WITHDRAW_ESCROW_THAWING_PERIOD + 1);
+
+        _withdrawEscrow(users.verifier, users.indexer);
+    }
+
     function testWithdraw_BalanceAfterCollect(
         uint256 amountDeposited,
         uint256 amountThawed,
         uint256 amountCollected
     ) public useGateway depositAndThawTokens(amountDeposited, amountThawed) {
         vm.assume(amountCollected > 0);
-        vm.assume(amountCollected <= amountDeposited);
+        vm.assume(amountCollected < amountDeposited);
 
         // burn some tokens to prevent overflow
         resetPrank(users.indexer);