[WIP] Horizon: fix unit tests and refactor

packages/contracts/contracts/upgrades/GraphProxy.sol

@@ -1,8 +1,6 @@
 // SPDX-License-Identifier: GPL-2.0-or-later
 
-pragma solidity ^0.7.6;
-
-import { Address } from "@openzeppelin/contracts/utils/Address.sol";
+pragma solidity >=0.6.12 <0.9.0;
 
 import { GraphProxyStorage } from "./GraphProxyStorage.sol";
 
@@ -160,7 +158,6 @@ contract GraphProxy is GraphProxyStorage, IGraphProxy {
      */
     function _acceptUpgrade() internal {
         address _pendingImplementation = _getPendingImplementation();
-        require(Address.isContract(_pendingImplementation), "Impl must be a contract");
         require(_pendingImplementation != address(0), "Impl cannot be zero address");
         require(msg.sender == _pendingImplementation, "Only pending implementation");
 

packages/contracts/contracts/upgrades/GraphProxyAdmin.sol

@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: GPL-2.0-or-later
 
-pragma solidity ^0.7.6;
+pragma solidity >=0.6.12 <0.9.0;
 
 import { Governed } from "../governance/Governed.sol";
 

packages/contracts/contracts/upgrades/GraphProxyStorage.sol

@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: GPL-2.0-or-later
 
-pragma solidity ^0.7.6;
+pragma solidity >=0.6.12 <0.9.0;
 
 /**
  * @title Graph Proxy Storage

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,22 +54,28 @@ 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 tokensReceiverRemaining = tokens - tokensProtocol - tokensDataService - tokensDelegationPool;
+        // Pay receiver
+        uint256 tokensReceiverRemaining = tokens - totalCut;
         _graphToken().pushTokens(receiver, tokensReceiverRemaining);
 
         emit GraphPaymentsCollected(

packages/horizon/contracts/payments/PaymentsEscrow.sol

@@ -78,14 +78,19 @@ contract PaymentsEscrow is Multicall, GraphDirectory, IPaymentsEscrow {
             revert GraphEscrowThawingPeriodTooLong(withdrawEscrowThawingPeriod, MAX_THAWING_PERIOD);
         }
 
-        revokeCollectorThawingPeriod = revokeCollectorThawingPeriod;
-        withdrawEscrowThawingPeriod = withdrawEscrowThawingPeriod;
+        REVOKE_COLLECTOR_THAWING_PERIOD = revokeCollectorThawingPeriod;
+        WITHDRAW_ESCROW_THAWING_PERIOD = withdrawEscrowThawingPeriod;
     }
 
     // approve a data service to collect funds
     function approveCollector(address dataService, uint256 amount) external {
-        authorizedCollectors[msg.sender][dataService].authorized = true;
-        authorizedCollectors[msg.sender][dataService].amount = amount;
+        Collector storage collector = authorizedCollectors[msg.sender][dataService];
+        if (collector.amount > amount) {
+            revert GraphEscrowCollectorInsufficientAmount(collector.amount, amount);
+        }
+
+        collector.authorized = true;
+        collector.amount = amount;
         emit AuthorizedCollector(msg.sender, dataService);
     }
 

packages/horizon/test/GraphBase.t.sol

@@ -3,6 +3,8 @@ pragma solidity ^0.8.24;
 
 import "forge-std/Test.sol";
 
+import { GraphProxyAdmin } from "@graphprotocol/contracts/contracts/upgrades/GraphProxyAdmin.sol";
+import { GraphProxy } from "@graphprotocol/contracts/contracts/upgrades/GraphProxy.sol";
 import { Controller } from "@graphprotocol/contracts/contracts/governance/Controller.sol";
 
 import { PaymentsEscrow } from "contracts/payments/PaymentsEscrow.sol";
@@ -18,6 +20,7 @@ abstract contract GraphBaseTest is Test, Constants {
 
     /* Contracts */
 
+    GraphProxyAdmin public proxyAdmin;
     Controller public controller;
     MockGRTToken public token;
     GraphPayments public payments;
@@ -41,6 +44,7 @@ abstract contract GraphBaseTest is Test, Constants {
 
     function setUp() public virtual {
         // Deploy ERC20 token
+        vm.prank(users.deployer);
         token = new MockGRTToken();
 
         // Setup Users
@@ -67,10 +71,16 @@ abstract contract GraphBaseTest is Test, Constants {
     }
 
     function deployProtocolContracts() private {
-        vm.prank(users.governor);
+        vm.startPrank(users.governor);
+        proxyAdmin = new GraphProxyAdmin();
         controller = new Controller();
+        vm.stopPrank();
 
-        // GraphPayments preddict address
+        // Staking Proxy
+        vm.prank(users.deployer);
+        GraphProxy stakingProxy = new GraphProxy(address(0), address(proxyAdmin));
+
+        // GraphPayments predict address
         bytes32 saltPayments = keccak256("GraphPaymentsSalt");
         bytes32 paymentsHash = keccak256(bytes.concat(
             vm.getCode("GraphPayments.sol:GraphPayments"),
@@ -82,7 +92,7 @@ abstract contract GraphBaseTest is Test, Constants {
             users.deployer
         );
 
-        // GraphEscrow preddict address
+        // GraphEscrow predict address
         bytes32 saltEscrow = keccak256("GraphEscrowSalt");
         bytes32 escrowHash = keccak256(bytes.concat(
             vm.getCode("GraphEscrow.sol:GraphEscrow"),
@@ -98,40 +108,18 @@ abstract contract GraphBaseTest is Test, Constants {
             users.deployer
         );
 
-        // HorizonStakingExtension preddict address
-        bytes32 saltHorizonStakingExtension = keccak256("HorizonStakingExtensionSalt");
-        bytes32 horizonStakingExtensionHash = keccak256(bytes.concat(
-            vm.getCode("HorizonStakingExtension.sol:HorizonStakingExtension"),
-            abi.encode(address(controller), subgraphDataServiceAddress)
-        ));
-        address predictedAddressHorizonStakingExtension = vm.computeCreate2Address(
-            saltHorizonStakingExtension,
-            horizonStakingExtensionHash,
-            users.deployer
-        );
-
-        // HorizonStaking preddict address
-        bytes32 saltHorizonStaking = keccak256("saltHorizonStaking");
-        bytes32 horizonStakingHash = keccak256(bytes.concat(
-            vm.getCode("HorizonStaking.sol:HorizonStaking"),
-            abi.encode(
-                address(controller), 
-                predictedAddressHorizonStakingExtension, 
-                subgraphDataServiceAddress
-            )
-        ));
-        // address predictedAddressHorizonStaking = vm.computeCreate2Address(
-        //     saltHorizonStaking,
-        //     horizonStakingHash,
-        //     users.deployer
-        // );
-
         // Setup controller
         vm.startPrank(users.governor);
         controller.setContractProxy(keccak256("GraphToken"), address(token));
-        controller.setContractProxy(keccak256("GraphEscrow"), predictedAddressEscrow);
+        controller.setContractProxy(keccak256("PaymentsEscrow"), predictedAddressEscrow);
         controller.setContractProxy(keccak256("GraphPayments"), predictedPaymentsAddress);
-        // controller.setContractProxy(keccak256("Staking"), predictedAddressHorizonStaking);
+        controller.setContractProxy(keccak256("Staking"), address(stakingProxy));
+        controller.setContractProxy(keccak256("EpochManager"), makeAddr("EpochManager"));
+        controller.setContractProxy(keccak256("RewardsManager"), makeAddr("RewardsManager"));
+        controller.setContractProxy(keccak256("Curation"), makeAddr("Curation"));
+        controller.setContractProxy(keccak256("GraphTokenGateway"), makeAddr("GraphTokenGateway"));
+        controller.setContractProxy(keccak256("BridgeEscrow"), makeAddr("BridgeEscrow"));
+        controller.setContractProxy(keccak256("GraphProxyAdmin"), makeAddr("GraphProxyAdmin"));
         vm.stopPrank();
 
         vm.startPrank(users.deployer);
@@ -144,16 +132,21 @@ abstract contract GraphBaseTest is Test, Constants {
             revokeCollectorThawingPeriod,
             withdrawEscrowThawingPeriod
         );
-        stakingBase = new HorizonStaking{salt: saltHorizonStaking}(
+        stakingExtension = new HorizonStakingExtension(
             address(controller),
-            predictedAddressHorizonStakingExtension,
             subgraphDataServiceAddress
         );
-        staking = IHorizonStaking(address(stakingBase));
-        // stakingExtension = new HorizonStakingExtension{salt: saltHorizonStakingExtension}(
-        //     address(controller),
-        //     subgraphDataServiceAddress
-        // );
+        stakingBase = new HorizonStaking(
+            address(controller),
+            address(stakingExtension),
+            subgraphDataServiceAddress
+        );
+        vm.stopPrank();
+
+        vm.startPrank(users.governor);
+        proxyAdmin.upgrade(stakingProxy, address(stakingBase));
+        proxyAdmin.acceptProxy(stakingBase, stakingProxy);
+        staking = IHorizonStaking(address(stakingProxy));
         vm.stopPrank();
     }
 

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

@@ -13,28 +13,36 @@ contract GraphEscrowTest is HorizonStakingSharedTest {
         vm.stopPrank();
     }
 
-    modifier approveEscrow(uint256 amount) {
-        _approveEscrow(amount);
+    modifier approveEscrow(uint256 tokens) {
+        changePrank(users.gateway);
+        _approveEscrow(tokens);
         _;
     }
 
-    modifier depositTokens(uint256 amount) {
-        vm.assume(amount > 0);
-        vm.assume(amount <= 10000 ether);
-        _depositTokens(amount);
+    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);
+        escrow.approveCollector(users.verifier, tokens);
         _;
     }
 
     function setUp() public virtual override {
         HorizonStakingSharedTest.setUp();
     }
 
-    function _depositTokens(uint256 amount) internal {
-        token.approve(address(escrow), amount);
-        escrow.deposit(users.indexer, amount);
+    function _depositTokens(uint256 tokens) internal {
+        token.approve(address(escrow), tokens);
+        escrow.deposit(users.indexer, tokens);
     }
 
-    function _approveEscrow(uint256 amount) internal {
-        token.approve(address(escrow), amount);
+    function _approveEscrow(uint256 tokens) internal {
+        token.approve(address(escrow), tokens);
     }
 }