Skip to content

Commit 971a5c5

Browse files
committed
chore: added unit tests for TAPCollector
1 parent 516ac3c commit 971a5c5

File tree

6 files changed

+223
-27
lines changed

6 files changed

+223
-27
lines changed

packages/horizon/test/GraphBase.t.sol

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import { TransparentUpgradeableProxy } from "@openzeppelin/contracts/proxy/trans
1111

1212
import { PaymentsEscrow } from "contracts/payments/PaymentsEscrow.sol";
1313
import { GraphPayments } from "contracts/payments/GraphPayments.sol";
14+
import { TAPCollector } from "contracts/payments/collectors/TAPCollector.sol";
1415
import { IHorizonStaking } from "contracts/interfaces/IHorizonStaking.sol";
1516
import { HorizonStaking } from "contracts/staking/HorizonStaking.sol";
1617
import { HorizonStakingExtension } from "contracts/staking/HorizonStakingExtension.sol";
@@ -39,6 +40,7 @@ abstract contract GraphBaseTest is IHorizonStakingTypes, Utils, Constants {
3940
EpochManagerMock public epochManager;
4041
RewardsManagerMock public rewardsManager;
4142
CurationMock public curation;
43+
TAPCollector tapCollector;
4244

4345
HorizonStaking private stakingBase;
4446
HorizonStakingExtension private stakingExtension;
@@ -84,6 +86,7 @@ abstract contract GraphBaseTest is IHorizonStakingTypes, Utils, Constants {
8486
vm.label({ account: address(escrow), newLabel: "PaymentsEscrow" });
8587
vm.label({ account: address(staking), newLabel: "HorizonStaking" });
8688
vm.label({ account: address(stakingExtension), newLabel: "HorizonStakingExtension" });
89+
vm.label({ account: address(tapCollector), newLabel: "TAPCollector" });
8790

8891
// Ensure caller is back to the original msg.sender
8992
vm.stopPrank();
@@ -182,6 +185,8 @@ abstract contract GraphBaseTest is IHorizonStakingTypes, Utils, Constants {
182185
subgraphDataServiceLegacyAddress
183186
);
184187

188+
tapCollector = new TAPCollector("TAPCollector", "1", address(controller));
189+
185190
resetPrank(users.governor);
186191
proxyAdmin.upgrade(stakingProxy, address(stakingBase));
187192
proxyAdmin.acceptProxy(stakingBase, stakingProxy);

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

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,9 @@ pragma solidity 0.8.27;
44
import "forge-std/Test.sol";
55

66
import { HorizonStakingSharedTest } from "../shared/horizon-staking/HorizonStakingShared.t.sol";
7+
import { PaymentsEscrowSharedTest } from "../shared/payments-escrow/PaymentsEscrowShared.t.sol";
78

8-
contract GraphEscrowTest is HorizonStakingSharedTest {
9+
contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
910

1011
/*
1112
* MODIFIERS
@@ -25,7 +26,7 @@ contract GraphEscrowTest is HorizonStakingSharedTest {
2526
modifier useDeposit(uint256 tokens) {
2627
vm.assume(tokens > 0);
2728
vm.assume(tokens <= MAX_STAKING_TOKENS);
28-
_depositTokens(tokens);
29+
_depositTokens(users.verifier, users.indexer, tokens);
2930
_;
3031
}
3132

@@ -38,7 +39,7 @@ contract GraphEscrowTest is HorizonStakingSharedTest {
3839
modifier depositAndThawTokens(uint256 amount, uint256 thawAmount) {
3940
vm.assume(thawAmount > 0);
4041
vm.assume(amount > thawAmount);
41-
_depositTokens(amount);
42+
_depositTokens(users.verifier, users.indexer, amount);
4243
escrow.thaw(users.verifier, users.indexer, thawAmount);
4344
_;
4445
}
@@ -47,11 +48,6 @@ contract GraphEscrowTest is HorizonStakingSharedTest {
4748
* HELPERS
4849
*/
4950

50-
function _depositTokens(uint256 tokens) internal {
51-
token.approve(address(escrow), tokens);
52-
escrow.deposit(users.verifier, users.indexer, tokens);
53-
}
54-
5551
function _approveEscrow(uint256 tokens) internal {
5652
token.approve(address(escrow), tokens);
5753
}

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,7 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
107107

108108
resetPrank(users.gateway);
109109
escrow.approveCollector(users.verifier, tokens);
110-
_depositTokens(tokens);
110+
_depositTokens(users.verifier, users.indexer, tokens);
111111

112112
resetPrank(users.verifier);
113113
_collect(IGraphPayments.PaymentTypes.QueryFee, users.gateway, users.indexer, tokens, subgraphDataServiceAddress, tokensDataService);
@@ -163,7 +163,7 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
163163

164164
resetPrank(users.gateway);
165165
escrow.approveCollector(users.verifier, amount);
166-
_depositTokens(amount);
166+
_depositTokens(users.verifier, users.indexer, amount);
167167

168168
resetPrank(users.verifier);
169169
vm.expectRevert(abi.encodeWithSelector(
@@ -182,7 +182,7 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
182182

183183
resetPrank(users.gateway);
184184
escrow.approveCollector(users.verifier, amount);
185-
_depositTokens(amount);
185+
_depositTokens(users.verifier, users.indexer, amount);
186186

187187
resetPrank(users.verifier);
188188
vm.expectRevert(abi.encodeWithSelector(

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

Lines changed: 1 addition & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -13,21 +13,6 @@ contract GraphEscrowCollectorTest is GraphEscrowTest {
1313
* HELPERS
1414
*/
1515

16-
function _approveCollector(uint256 tokens) internal {
17-
(uint256 beforeAllowance,) = escrow.authorizedCollectors(users.gateway, users.verifier);
18-
vm.expectEmit(address(escrow));
19-
emit IPaymentsEscrow.AuthorizedCollector(
20-
users.gateway, // payer
21-
users.verifier, // collector
22-
tokens, // addedAllowance
23-
beforeAllowance + tokens // newTotalAllowance after the added allowance
24-
);
25-
escrow.approveCollector(users.verifier, tokens);
26-
(uint256 allowance, uint256 thawEndTimestamp) = escrow.authorizedCollectors(users.gateway, users.verifier);
27-
assertEq(allowance - beforeAllowance, tokens);
28-
assertEq(thawEndTimestamp, 0);
29-
}
30-
3116
function _thawCollector() internal {
3217
(uint256 beforeAllowance,) = escrow.authorizedCollectors(users.gateway, users.verifier);
3318
vm.expectEmit(address(escrow));
@@ -74,7 +59,7 @@ contract GraphEscrowCollectorTest is GraphEscrowTest {
7459

7560
uint256 approveTokens = tokens / approveSteps;
7661
for (uint i = 0; i < approveSteps; i++) {
77-
_approveCollector(approveTokens);
62+
_approveCollector(users.verifier, approveTokens);
7863
}
7964
}
8065

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
// SPDX-License-Identifier: MIT
2+
pragma solidity 0.8.27;
3+
4+
import "forge-std/Test.sol";
5+
6+
import { ECDSA } from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";
7+
import { IHorizonStakingMain } from "../../contracts/interfaces/internal/IHorizonStakingMain.sol";
8+
import { ITAPCollector } from "../../contracts/interfaces/ITAPCollector.sol";
9+
import { IPaymentsCollector } from "../../contracts/interfaces/IPaymentsCollector.sol";
10+
import { IGraphPayments } from "../../contracts/interfaces/IGraphPayments.sol";
11+
import { TAPCollector } from "../../contracts/payments/collectors/TAPCollector.sol";
12+
import { PPMMath } from "../../contracts/libraries/PPMMath.sol";
13+
14+
import { HorizonStakingSharedTest } from "../shared/horizon-staking/HorizonStakingShared.t.sol";
15+
import { PaymentsEscrowSharedTest } from "../shared/payments-escrow/PaymentsEscrowShared.t.sol";
16+
17+
contract TAPCollectorTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
18+
using PPMMath for uint256;
19+
20+
address payer;
21+
uint256 payerPrivateKey;
22+
23+
/*
24+
* HELPERS
25+
*/
26+
27+
function _getQueryFeeEncodedData(address indexer, address collector, uint128 tokens) private view returns (bytes memory) {
28+
ITAPCollector.ReceiptAggregateVoucher memory rav = _getRAV(indexer, collector, tokens);
29+
bytes32 messageHash = tapCollector.encodeRAV(rav);
30+
(uint8 v, bytes32 r, bytes32 s) = vm.sign(payerPrivateKey, messageHash);
31+
bytes memory signature = abi.encodePacked(r, s, v);
32+
ITAPCollector.SignedRAV memory signedRAV = ITAPCollector.SignedRAV(rav, signature);
33+
return abi.encode(signedRAV);
34+
}
35+
36+
function _getRAV(
37+
address indexer,
38+
address collector,
39+
uint128 tokens
40+
) private pure returns (ITAPCollector.ReceiptAggregateVoucher memory rav) {
41+
return
42+
ITAPCollector.ReceiptAggregateVoucher({
43+
dataService: collector,
44+
serviceProvider: indexer,
45+
timestampNs: 0,
46+
valueAggregate: tokens,
47+
metadata: abi.encode("")
48+
});
49+
}
50+
51+
function _collect(IGraphPayments.PaymentTypes _paymentType, bytes memory _data) private {
52+
(ITAPCollector.SignedRAV memory signedRAV, uint256 dataServiceCut) = abi.decode(_data, (ITAPCollector.SignedRAV, uint256));
53+
bytes32 messageHash = tapCollector.encodeRAV(signedRAV.rav);
54+
address _payer = ECDSA.recover(messageHash, signedRAV.signature);
55+
uint256 tokensAlreadyCollected = tapCollector.tokensCollected(signedRAV.rav.dataService, signedRAV.rav.serviceProvider, _payer);
56+
uint256 tokensToCollect = signedRAV.rav.valueAggregate - tokensAlreadyCollected;
57+
uint256 tokensDataService = tokensToCollect.mulPPM(dataServiceCut);
58+
59+
vm.expectEmit(address(tapCollector));
60+
emit IPaymentsCollector.PaymentCollected(
61+
_paymentType,
62+
_payer,
63+
signedRAV.rav.serviceProvider,
64+
tokensToCollect,
65+
signedRAV.rav.dataService,
66+
tokensDataService
67+
);
68+
emit ITAPCollector.RAVCollected(
69+
_payer,
70+
signedRAV.rav.dataService,
71+
signedRAV.rav.serviceProvider,
72+
signedRAV.rav.timestampNs,
73+
signedRAV.rav.valueAggregate,
74+
signedRAV.rav.metadata,
75+
signedRAV.signature
76+
);
77+
78+
uint256 tokensCollected = tapCollector.collect(_paymentType, _data);
79+
assertEq(tokensCollected, tokensToCollect);
80+
81+
uint256 tokensCollectedAfter = tapCollector.tokensCollected(signedRAV.rav.dataService, signedRAV.rav.serviceProvider, _payer);
82+
assertEq(tokensCollectedAfter, signedRAV.rav.valueAggregate);
83+
}
84+
85+
/*
86+
* SET UP
87+
*/
88+
89+
function setUp() public virtual override {
90+
super.setUp();
91+
(payer, payerPrivateKey) = makeAddrAndKey("payer");
92+
vm.label({ account: payer, newLabel: "payer" });
93+
deal({ token: address(token), to: payer, give: type(uint256).max });
94+
}
95+
96+
/*
97+
* TESTS
98+
*/
99+
100+
function testCollect(uint256 tokens) public {
101+
tokens = bound(tokens, 1, type(uint128).max);
102+
103+
resetPrank(payer);
104+
_approveCollector(address(tapCollector), tokens);
105+
_depositTokens(address(tapCollector), users.indexer, tokens);
106+
bytes memory data = _getQueryFeeEncodedData(users.indexer, users.verifier, uint128(tokens));
107+
108+
resetPrank(users.verifier);
109+
_collect(IGraphPayments.PaymentTypes.QueryFee, data);
110+
}
111+
112+
function testCollect_Multiple(uint256 tokens, uint8 steps) public {
113+
steps = uint8(bound(steps, 1, 100));
114+
tokens = bound(tokens, steps, type(uint128).max);
115+
116+
resetPrank(payer);
117+
_approveCollector(address(tapCollector), tokens);
118+
_depositTokens(address(tapCollector), users.indexer, tokens);
119+
120+
resetPrank(users.verifier);
121+
uint256 payed = 0;
122+
uint256 tokensPerStep = tokens / steps;
123+
for (uint256 i = 0; i < steps; i++) {
124+
bytes memory data = _getQueryFeeEncodedData(users.indexer, users.verifier, uint128(payed + tokensPerStep));
125+
_collect(IGraphPayments.PaymentTypes.QueryFee, data);
126+
payed += tokensPerStep;
127+
}
128+
}
129+
130+
function testCollect_RevertWhen_CallerNotDataService(uint256 tokens) public {
131+
tokens = bound(tokens, 1, type(uint128).max);
132+
133+
resetPrank(payer);
134+
_approveCollector(address(tapCollector), tokens);
135+
_depositTokens(address(tapCollector), users.indexer, tokens);
136+
bytes memory data = _getQueryFeeEncodedData(users.indexer, users.verifier, uint128(tokens));
137+
138+
resetPrank(users.indexer);
139+
bytes memory expectedError = abi.encodeWithSelector(
140+
ITAPCollector.TAPCollectorCallerNotDataService.selector,
141+
users.indexer,
142+
users.verifier
143+
);
144+
vm.expectRevert(expectedError);
145+
tapCollector.collect(IGraphPayments.PaymentTypes.QueryFee, data);
146+
}
147+
148+
function testCollect_RevertWhen_InconsistentRAVTokens(uint256 tokens) public {
149+
tokens = bound(tokens, 1, type(uint128).max);
150+
151+
resetPrank(payer);
152+
_approveCollector(address(tapCollector), tokens);
153+
_depositTokens(address(tapCollector), users.indexer, tokens);
154+
bytes memory data = _getQueryFeeEncodedData(users.indexer, users.verifier, uint128(tokens));
155+
156+
resetPrank(users.verifier);
157+
_collect(IGraphPayments.PaymentTypes.QueryFee, data);
158+
159+
// Attempt to collect again
160+
vm.expectRevert(abi.encodeWithSelector(
161+
ITAPCollector.TAPCollectorInconsistentRAVTokens.selector,
162+
tokens,
163+
tokens
164+
));
165+
tapCollector.collect(IGraphPayments.PaymentTypes.QueryFee, data);
166+
}
167+
}
Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
// SPDX-License-Identifier: MIT
2+
pragma solidity 0.8.27;
3+
4+
import "forge-std/Test.sol";
5+
6+
import { IPaymentsEscrow } from "../../../contracts/interfaces/IPaymentsEscrow.sol";
7+
import { GraphBaseTest } from "../../GraphBase.t.sol";
8+
9+
abstract contract PaymentsEscrowSharedTest is GraphBaseTest {
10+
11+
/*
12+
* HELPERS
13+
*/
14+
15+
function _approveCollector(address _verifier, uint256 _tokens) internal {
16+
(, address msgSender, ) = vm.readCallers();
17+
(uint256 beforeAllowance,) = escrow.authorizedCollectors(msgSender, _verifier);
18+
vm.expectEmit(address(escrow));
19+
emit IPaymentsEscrow.AuthorizedCollector(
20+
msgSender, // payer
21+
_verifier, // collector
22+
_tokens, // addedAllowance
23+
beforeAllowance + _tokens // newTotalAllowance after the added allowance
24+
);
25+
escrow.approveCollector(_verifier, _tokens);
26+
(uint256 allowance, uint256 thawEndTimestamp) = escrow.authorizedCollectors(msgSender, _verifier);
27+
assertEq(allowance - beforeAllowance, _tokens);
28+
assertEq(thawEndTimestamp, 0);
29+
}
30+
31+
function _depositTokens(address _collector, address _receiver, uint256 _tokens) internal {
32+
(, address msgSender, ) = vm.readCallers();
33+
(uint256 escrowBalanceBefore,,) = escrow.escrowAccounts(msgSender, _collector, _receiver);
34+
token.approve(address(escrow), _tokens);
35+
36+
vm.expectEmit(address(escrow));
37+
emit IPaymentsEscrow.Deposit(msgSender, _collector, _receiver, _tokens);
38+
escrow.deposit(_collector, _receiver, _tokens);
39+
40+
(uint256 escrowBalanceAfter,,) = escrow.escrowAccounts(msgSender, _collector, _receiver);
41+
assertEq(escrowBalanceAfter - _tokens, escrowBalanceBefore);
42+
}
43+
}

0 commit comments

Comments
 (0)