From 971a5c5afd1a6caa88568932c68799cda3c87cab Mon Sep 17 00:00:00 2001 From: Miguel de Elias Date: Thu, 10 Oct 2024 14:02:46 -0300 Subject: [PATCH] chore: added unit tests for TAPCollector --- packages/horizon/test/GraphBase.t.sol | 5 + .../horizon/test/escrow/GraphEscrow.t.sol | 12 +- packages/horizon/test/escrow/collect.t.sol | 6 +- packages/horizon/test/escrow/collector.t.sol | 17 +- .../horizon/test/payments/TAPCollector.t.sol | 167 ++++++++++++++++++ .../PaymentsEscrowShared.t.sol | 43 +++++ 6 files changed, 223 insertions(+), 27 deletions(-) create mode 100644 packages/horizon/test/payments/TAPCollector.t.sol create mode 100644 packages/horizon/test/shared/payments-escrow/PaymentsEscrowShared.t.sol diff --git a/packages/horizon/test/GraphBase.t.sol b/packages/horizon/test/GraphBase.t.sol index f02733525..16c89e5c5 100644 --- a/packages/horizon/test/GraphBase.t.sol +++ b/packages/horizon/test/GraphBase.t.sol @@ -11,6 +11,7 @@ import { TransparentUpgradeableProxy } from "@openzeppelin/contracts/proxy/trans import { PaymentsEscrow } from "contracts/payments/PaymentsEscrow.sol"; import { GraphPayments } from "contracts/payments/GraphPayments.sol"; +import { TAPCollector } from "contracts/payments/collectors/TAPCollector.sol"; import { IHorizonStaking } from "contracts/interfaces/IHorizonStaking.sol"; import { HorizonStaking } from "contracts/staking/HorizonStaking.sol"; import { HorizonStakingExtension } from "contracts/staking/HorizonStakingExtension.sol"; @@ -39,6 +40,7 @@ abstract contract GraphBaseTest is IHorizonStakingTypes, Utils, Constants { EpochManagerMock public epochManager; RewardsManagerMock public rewardsManager; CurationMock public curation; + TAPCollector tapCollector; HorizonStaking private stakingBase; HorizonStakingExtension private stakingExtension; @@ -84,6 +86,7 @@ abstract contract GraphBaseTest is IHorizonStakingTypes, Utils, Constants { vm.label({ account: address(escrow), newLabel: "PaymentsEscrow" }); vm.label({ account: address(staking), newLabel: "HorizonStaking" }); vm.label({ account: address(stakingExtension), newLabel: "HorizonStakingExtension" }); + vm.label({ account: address(tapCollector), newLabel: "TAPCollector" }); // Ensure caller is back to the original msg.sender vm.stopPrank(); @@ -182,6 +185,8 @@ abstract contract GraphBaseTest is IHorizonStakingTypes, Utils, Constants { subgraphDataServiceLegacyAddress ); + tapCollector = new TAPCollector("TAPCollector", "1", address(controller)); + resetPrank(users.governor); proxyAdmin.upgrade(stakingProxy, address(stakingBase)); proxyAdmin.acceptProxy(stakingBase, stakingProxy); diff --git a/packages/horizon/test/escrow/GraphEscrow.t.sol b/packages/horizon/test/escrow/GraphEscrow.t.sol index 75d0b55ae..a472162e9 100644 --- a/packages/horizon/test/escrow/GraphEscrow.t.sol +++ b/packages/horizon/test/escrow/GraphEscrow.t.sol @@ -4,8 +4,9 @@ pragma solidity 0.8.27; import "forge-std/Test.sol"; import { HorizonStakingSharedTest } from "../shared/horizon-staking/HorizonStakingShared.t.sol"; +import { PaymentsEscrowSharedTest } from "../shared/payments-escrow/PaymentsEscrowShared.t.sol"; -contract GraphEscrowTest is HorizonStakingSharedTest { +contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest { /* * MODIFIERS @@ -25,7 +26,7 @@ contract GraphEscrowTest is HorizonStakingSharedTest { modifier useDeposit(uint256 tokens) { vm.assume(tokens > 0); vm.assume(tokens <= MAX_STAKING_TOKENS); - _depositTokens(tokens); + _depositTokens(users.verifier, users.indexer, tokens); _; } @@ -38,7 +39,7 @@ contract GraphEscrowTest is HorizonStakingSharedTest { modifier depositAndThawTokens(uint256 amount, uint256 thawAmount) { vm.assume(thawAmount > 0); vm.assume(amount > thawAmount); - _depositTokens(amount); + _depositTokens(users.verifier, users.indexer, amount); escrow.thaw(users.verifier, users.indexer, thawAmount); _; } @@ -47,11 +48,6 @@ contract GraphEscrowTest is HorizonStakingSharedTest { * HELPERS */ - function _depositTokens(uint256 tokens) internal { - token.approve(address(escrow), tokens); - escrow.deposit(users.verifier, users.indexer, tokens); - } - function _approveEscrow(uint256 tokens) internal { token.approve(address(escrow), tokens); } diff --git a/packages/horizon/test/escrow/collect.t.sol b/packages/horizon/test/escrow/collect.t.sol index db0c2c099..72b795ee9 100644 --- a/packages/horizon/test/escrow/collect.t.sol +++ b/packages/horizon/test/escrow/collect.t.sol @@ -107,7 +107,7 @@ contract GraphEscrowCollectTest is GraphEscrowTest { resetPrank(users.gateway); escrow.approveCollector(users.verifier, tokens); - _depositTokens(tokens); + _depositTokens(users.verifier, users.indexer, tokens); resetPrank(users.verifier); _collect(IGraphPayments.PaymentTypes.QueryFee, users.gateway, users.indexer, tokens, subgraphDataServiceAddress, tokensDataService); @@ -163,7 +163,7 @@ contract GraphEscrowCollectTest is GraphEscrowTest { resetPrank(users.gateway); escrow.approveCollector(users.verifier, amount); - _depositTokens(amount); + _depositTokens(users.verifier, users.indexer, amount); resetPrank(users.verifier); vm.expectRevert(abi.encodeWithSelector( @@ -182,7 +182,7 @@ contract GraphEscrowCollectTest is GraphEscrowTest { resetPrank(users.gateway); escrow.approveCollector(users.verifier, amount); - _depositTokens(amount); + _depositTokens(users.verifier, users.indexer, amount); resetPrank(users.verifier); vm.expectRevert(abi.encodeWithSelector( diff --git a/packages/horizon/test/escrow/collector.t.sol b/packages/horizon/test/escrow/collector.t.sol index ed5c05384..d6cb3bc0f 100644 --- a/packages/horizon/test/escrow/collector.t.sol +++ b/packages/horizon/test/escrow/collector.t.sol @@ -13,21 +13,6 @@ contract GraphEscrowCollectorTest is GraphEscrowTest { * HELPERS */ - function _approveCollector(uint256 tokens) internal { - (uint256 beforeAllowance,) = escrow.authorizedCollectors(users.gateway, users.verifier); - vm.expectEmit(address(escrow)); - emit IPaymentsEscrow.AuthorizedCollector( - users.gateway, // payer - users.verifier, // collector - tokens, // addedAllowance - beforeAllowance + tokens // newTotalAllowance after the added allowance - ); - escrow.approveCollector(users.verifier, tokens); - (uint256 allowance, uint256 thawEndTimestamp) = escrow.authorizedCollectors(users.gateway, users.verifier); - assertEq(allowance - beforeAllowance, tokens); - assertEq(thawEndTimestamp, 0); - } - function _thawCollector() internal { (uint256 beforeAllowance,) = escrow.authorizedCollectors(users.gateway, users.verifier); vm.expectEmit(address(escrow)); @@ -74,7 +59,7 @@ contract GraphEscrowCollectorTest is GraphEscrowTest { uint256 approveTokens = tokens / approveSteps; for (uint i = 0; i < approveSteps; i++) { - _approveCollector(approveTokens); + _approveCollector(users.verifier, approveTokens); } } diff --git a/packages/horizon/test/payments/TAPCollector.t.sol b/packages/horizon/test/payments/TAPCollector.t.sol new file mode 100644 index 000000000..e1c177282 --- /dev/null +++ b/packages/horizon/test/payments/TAPCollector.t.sol @@ -0,0 +1,167 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.27; + +import "forge-std/Test.sol"; + +import { ECDSA } from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol"; +import { IHorizonStakingMain } from "../../contracts/interfaces/internal/IHorizonStakingMain.sol"; +import { ITAPCollector } from "../../contracts/interfaces/ITAPCollector.sol"; +import { IPaymentsCollector } from "../../contracts/interfaces/IPaymentsCollector.sol"; +import { IGraphPayments } from "../../contracts/interfaces/IGraphPayments.sol"; +import { TAPCollector } from "../../contracts/payments/collectors/TAPCollector.sol"; +import { PPMMath } from "../../contracts/libraries/PPMMath.sol"; + +import { HorizonStakingSharedTest } from "../shared/horizon-staking/HorizonStakingShared.t.sol"; +import { PaymentsEscrowSharedTest } from "../shared/payments-escrow/PaymentsEscrowShared.t.sol"; + +contract TAPCollectorTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest { + using PPMMath for uint256; + + address payer; + uint256 payerPrivateKey; + + /* + * HELPERS + */ + + function _getQueryFeeEncodedData(address indexer, address collector, uint128 tokens) private view returns (bytes memory) { + ITAPCollector.ReceiptAggregateVoucher memory rav = _getRAV(indexer, collector, tokens); + bytes32 messageHash = tapCollector.encodeRAV(rav); + (uint8 v, bytes32 r, bytes32 s) = vm.sign(payerPrivateKey, messageHash); + bytes memory signature = abi.encodePacked(r, s, v); + ITAPCollector.SignedRAV memory signedRAV = ITAPCollector.SignedRAV(rav, signature); + return abi.encode(signedRAV); + } + + function _getRAV( + address indexer, + address collector, + uint128 tokens + ) private pure returns (ITAPCollector.ReceiptAggregateVoucher memory rav) { + return + ITAPCollector.ReceiptAggregateVoucher({ + dataService: collector, + serviceProvider: indexer, + timestampNs: 0, + valueAggregate: tokens, + metadata: abi.encode("") + }); + } + + function _collect(IGraphPayments.PaymentTypes _paymentType, bytes memory _data) private { + (ITAPCollector.SignedRAV memory signedRAV, uint256 dataServiceCut) = abi.decode(_data, (ITAPCollector.SignedRAV, uint256)); + bytes32 messageHash = tapCollector.encodeRAV(signedRAV.rav); + address _payer = ECDSA.recover(messageHash, signedRAV.signature); + uint256 tokensAlreadyCollected = tapCollector.tokensCollected(signedRAV.rav.dataService, signedRAV.rav.serviceProvider, _payer); + uint256 tokensToCollect = signedRAV.rav.valueAggregate - tokensAlreadyCollected; + uint256 tokensDataService = tokensToCollect.mulPPM(dataServiceCut); + + vm.expectEmit(address(tapCollector)); + emit IPaymentsCollector.PaymentCollected( + _paymentType, + _payer, + signedRAV.rav.serviceProvider, + tokensToCollect, + signedRAV.rav.dataService, + tokensDataService + ); + emit ITAPCollector.RAVCollected( + _payer, + signedRAV.rav.dataService, + signedRAV.rav.serviceProvider, + signedRAV.rav.timestampNs, + signedRAV.rav.valueAggregate, + signedRAV.rav.metadata, + signedRAV.signature + ); + + uint256 tokensCollected = tapCollector.collect(_paymentType, _data); + assertEq(tokensCollected, tokensToCollect); + + uint256 tokensCollectedAfter = tapCollector.tokensCollected(signedRAV.rav.dataService, signedRAV.rav.serviceProvider, _payer); + assertEq(tokensCollectedAfter, signedRAV.rav.valueAggregate); + } + + /* + * SET UP + */ + + function setUp() public virtual override { + super.setUp(); + (payer, payerPrivateKey) = makeAddrAndKey("payer"); + vm.label({ account: payer, newLabel: "payer" }); + deal({ token: address(token), to: payer, give: type(uint256).max }); + } + + /* + * TESTS + */ + + function testCollect(uint256 tokens) public { + tokens = bound(tokens, 1, type(uint128).max); + + resetPrank(payer); + _approveCollector(address(tapCollector), tokens); + _depositTokens(address(tapCollector), users.indexer, tokens); + bytes memory data = _getQueryFeeEncodedData(users.indexer, users.verifier, uint128(tokens)); + + resetPrank(users.verifier); + _collect(IGraphPayments.PaymentTypes.QueryFee, data); + } + + function testCollect_Multiple(uint256 tokens, uint8 steps) public { + steps = uint8(bound(steps, 1, 100)); + tokens = bound(tokens, steps, type(uint128).max); + + resetPrank(payer); + _approveCollector(address(tapCollector), tokens); + _depositTokens(address(tapCollector), users.indexer, tokens); + + resetPrank(users.verifier); + uint256 payed = 0; + uint256 tokensPerStep = tokens / steps; + for (uint256 i = 0; i < steps; i++) { + bytes memory data = _getQueryFeeEncodedData(users.indexer, users.verifier, uint128(payed + tokensPerStep)); + _collect(IGraphPayments.PaymentTypes.QueryFee, data); + payed += tokensPerStep; + } + } + + function testCollect_RevertWhen_CallerNotDataService(uint256 tokens) public { + tokens = bound(tokens, 1, type(uint128).max); + + resetPrank(payer); + _approveCollector(address(tapCollector), tokens); + _depositTokens(address(tapCollector), users.indexer, tokens); + bytes memory data = _getQueryFeeEncodedData(users.indexer, users.verifier, uint128(tokens)); + + resetPrank(users.indexer); + bytes memory expectedError = abi.encodeWithSelector( + ITAPCollector.TAPCollectorCallerNotDataService.selector, + users.indexer, + users.verifier + ); + vm.expectRevert(expectedError); + tapCollector.collect(IGraphPayments.PaymentTypes.QueryFee, data); + } + + function testCollect_RevertWhen_InconsistentRAVTokens(uint256 tokens) public { + tokens = bound(tokens, 1, type(uint128).max); + + resetPrank(payer); + _approveCollector(address(tapCollector), tokens); + _depositTokens(address(tapCollector), users.indexer, tokens); + bytes memory data = _getQueryFeeEncodedData(users.indexer, users.verifier, uint128(tokens)); + + resetPrank(users.verifier); + _collect(IGraphPayments.PaymentTypes.QueryFee, data); + + // Attempt to collect again + vm.expectRevert(abi.encodeWithSelector( + ITAPCollector.TAPCollectorInconsistentRAVTokens.selector, + tokens, + tokens + )); + tapCollector.collect(IGraphPayments.PaymentTypes.QueryFee, data); + } +} diff --git a/packages/horizon/test/shared/payments-escrow/PaymentsEscrowShared.t.sol b/packages/horizon/test/shared/payments-escrow/PaymentsEscrowShared.t.sol new file mode 100644 index 000000000..b7ee76839 --- /dev/null +++ b/packages/horizon/test/shared/payments-escrow/PaymentsEscrowShared.t.sol @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.27; + +import "forge-std/Test.sol"; + +import { IPaymentsEscrow } from "../../../contracts/interfaces/IPaymentsEscrow.sol"; +import { GraphBaseTest } from "../../GraphBase.t.sol"; + +abstract contract PaymentsEscrowSharedTest is GraphBaseTest { + + /* + * HELPERS + */ + + function _approveCollector(address _verifier, uint256 _tokens) internal { + (, address msgSender, ) = vm.readCallers(); + (uint256 beforeAllowance,) = escrow.authorizedCollectors(msgSender, _verifier); + vm.expectEmit(address(escrow)); + emit IPaymentsEscrow.AuthorizedCollector( + msgSender, // payer + _verifier, // collector + _tokens, // addedAllowance + beforeAllowance + _tokens // newTotalAllowance after the added allowance + ); + escrow.approveCollector(_verifier, _tokens); + (uint256 allowance, uint256 thawEndTimestamp) = escrow.authorizedCollectors(msgSender, _verifier); + assertEq(allowance - beforeAllowance, _tokens); + assertEq(thawEndTimestamp, 0); + } + + function _depositTokens(address _collector, address _receiver, uint256 _tokens) internal { + (, address msgSender, ) = vm.readCallers(); + (uint256 escrowBalanceBefore,,) = escrow.escrowAccounts(msgSender, _collector, _receiver); + token.approve(address(escrow), _tokens); + + vm.expectEmit(address(escrow)); + emit IPaymentsEscrow.Deposit(msgSender, _collector, _receiver, _tokens); + escrow.deposit(_collector, _receiver, _tokens); + + (uint256 escrowBalanceAfter,,) = escrow.escrowAccounts(msgSender, _collector, _receiver); + assertEq(escrowBalanceAfter - _tokens, escrowBalanceBefore); + } +}