@@ -113,13 +113,21 @@ contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
113113 uint256 payerEscrowBalance;
114114 }
115115
116+ struct CollectTokensData {
117+ uint256 tokensProtocol;
118+ uint256 tokensDataService;
119+ uint256 tokensDelegation;
120+ uint256 receiverExpectedPayment;
121+ }
122+
116123 function _collectEscrow (
117124 IGraphPayments.PaymentTypes _paymentType ,
118125 address _payer ,
119126 address _receiver ,
120127 uint256 _tokens ,
121128 address _dataService ,
122- uint256 _dataServiceCut
129+ uint256 _dataServiceCut ,
130+ address _paymentsDestination
123131 ) internal {
124132 (, address _collector , ) = vm.readCallers ();
125133
@@ -132,31 +140,49 @@ contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
132140 dataServiceBalance: token.balanceOf (_dataService),
133141 payerEscrowBalance: 0
134142 });
143+ CollectTokensData memory collectTokensData = CollectTokensData ({
144+ tokensProtocol: 0 ,
145+ tokensDataService: 0 ,
146+ tokensDelegation: 0 ,
147+ receiverExpectedPayment: 0
148+ });
135149
136150 {
137151 (uint256 payerEscrowBalance , , ) = escrow.escrowAccounts (_payer, _collector, _receiver);
138152 previousBalances.payerEscrowBalance = payerEscrowBalance;
139153 }
140154
141155 vm.expectEmit (address (escrow));
142- emit IPaymentsEscrow.EscrowCollected (_paymentType, _payer, _collector, _receiver, _tokens);
143- escrow.collect (_paymentType, _payer, _receiver, _tokens, _dataService, _dataServiceCut);
156+ emit IPaymentsEscrow.EscrowCollected (
157+ _paymentType,
158+ _payer,
159+ _collector,
160+ _receiver,
161+ _tokens,
162+ _paymentsDestination
163+ );
164+ escrow.collect (_paymentType, _payer, _receiver, _tokens, _dataService, _dataServiceCut, _paymentsDestination);
144165
145166 // Calculate cuts
146167 // this is nasty but stack is indeed too deep
147- uint256 tokensDataService = (_tokens - _tokens.mulPPMRoundUp (payments.PROTOCOL_PAYMENT_CUT ())).mulPPMRoundUp (
168+ collectTokensData.tokensProtocol = _tokens.mulPPMRoundUp (payments.PROTOCOL_PAYMENT_CUT ());
169+ collectTokensData.tokensDataService = (_tokens - collectTokensData.tokensProtocol).mulPPMRoundUp (
148170 _dataServiceCut
149171 );
150- uint256 tokensDelegation = 0 ;
172+
151173 IHorizonStakingTypes.DelegationPool memory pool = staking.getDelegationPool (_receiver, _dataService);
152174 if (pool.shares > 0 ) {
153- tokensDelegation = (_tokens - _tokens.mulPPMRoundUp (payments.PROTOCOL_PAYMENT_CUT ()) - tokensDataService)
154- .mulPPMRoundUp (staking.getDelegationFeeCut (_receiver, _dataService, _paymentType));
175+ collectTokensData.tokensDelegation = (_tokens -
176+ collectTokensData.tokensProtocol -
177+ collectTokensData.tokensDataService).mulPPMRoundUp (
178+ staking.getDelegationFeeCut (_receiver, _dataService, _paymentType)
179+ );
155180 }
156- uint256 receiverExpectedPayment = _tokens -
157- _tokens.mulPPMRoundUp (payments.PROTOCOL_PAYMENT_CUT ()) -
158- tokensDataService -
159- tokensDelegation;
181+ collectTokensData.receiverExpectedPayment =
182+ _tokens -
183+ collectTokensData.tokensProtocol -
184+ collectTokensData.tokensDataService -
185+ collectTokensData.tokensDelegation;
160186
161187 // After balances
162188 CollectPaymentData memory afterBalances = CollectPaymentData ({
@@ -173,11 +199,11 @@ contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
173199 }
174200
175201 // Check receiver balance after payment
176- assertEq (afterBalances.receiverBalance - previousBalances.receiverBalance, receiverExpectedPayment);
202+ assertEq (afterBalances.receiverBalance - previousBalances.receiverBalance, collectTokensData. receiverExpectedPayment);
177203 assertEq (token.balanceOf (address (payments)), 0 );
178204
179205 // Check delegation pool balance after payment
180- assertEq (afterBalances.delegationPoolBalance - previousBalances.delegationPoolBalance, tokensDelegation);
206+ assertEq (afterBalances.delegationPoolBalance - previousBalances.delegationPoolBalance, collectTokensData. tokensDelegation);
181207
182208 // Check that the escrow account has been updated
183209 assertEq (previousBalances.escrowBalance, afterBalances.escrowBalance + _tokens);
@@ -186,7 +212,7 @@ contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
186212 assertEq (previousBalances.paymentsBalance, afterBalances.paymentsBalance);
187213
188214 // Check data service balance after payment
189- assertEq (afterBalances.dataServiceBalance - previousBalances.dataServiceBalance, tokensDataService);
215+ assertEq (afterBalances.dataServiceBalance - previousBalances.dataServiceBalance, collectTokensData. tokensDataService);
190216
191217 // Check payers escrow balance after payment
192218 assertEq (previousBalances.payerEscrowBalance - _tokens, afterBalances.payerEscrowBalance);
0 commit comments