@@ -30,7 +30,9 @@ contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
3030 }
3131
3232 modifier depositAndThawTokens (uint256 amount , uint256 thawAmount ) {
33+ vm.assume (amount > 0 );
3334 vm.assume (thawAmount > 0 );
35+ vm.assume (amount <= MAX_STAKING_TOKENS);
3436 vm.assume (amount > thawAmount);
3537 _depositTokens (users.verifier, users.indexer, amount);
3638 escrow.thaw (users.verifier, users.indexer, thawAmount);
@@ -62,14 +64,38 @@ contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
6264 (, uint256 amountThawingBefore , uint256 thawEndTimestampBefore ) = escrow.escrowAccounts (msgSender, collector, receiver);
6365
6466 vm.expectEmit (address (escrow));
65- emit IPaymentsEscrow.CancelThaw (msgSender, receiver, amountThawingBefore, thawEndTimestampBefore);
67+ emit IPaymentsEscrow.CancelThaw (msgSender, collector, receiver, amountThawingBefore, thawEndTimestampBefore);
6668 escrow.cancelThaw (collector, receiver);
6769
6870 (, uint256 amountThawing , uint256 thawEndTimestamp ) = escrow.escrowAccounts (msgSender, collector, receiver);
6971 assertEq (amountThawing, 0 );
7072 assertEq (thawEndTimestamp, 0 );
7173 }
7274
75+ function _withdrawEscrow (address collector , address receiver ) internal {
76+ (, address msgSender , ) = vm.readCallers ();
77+
78+ (uint256 balanceBefore , uint256 amountThawingBefore , ) = escrow.escrowAccounts (msgSender, collector, receiver);
79+ uint256 tokenBalanceBeforeSender = token.balanceOf (msgSender);
80+ uint256 tokenBalanceBeforeEscrow = token.balanceOf (address (escrow));
81+
82+ uint256 amountToWithdraw = amountThawingBefore > balanceBefore ? balanceBefore : amountThawingBefore;
83+ vm.expectEmit (address (escrow));
84+ emit IPaymentsEscrow.Withdraw (msgSender, collector, receiver, amountToWithdraw);
85+ escrow.withdraw (collector, receiver);
86+
87+ (uint256 balanceAfter , uint256 tokensThawingAfter , uint256 thawEndTimestampAfter ) = escrow.escrowAccounts (msgSender, collector, receiver);
88+ uint256 tokenBalanceAfterSender = token.balanceOf (msgSender);
89+ uint256 tokenBalanceAfterEscrow = token.balanceOf (address (escrow));
90+
91+ assertEq (balanceAfter, balanceBefore - amountToWithdraw);
92+ assertEq (tokensThawingAfter, 0 );
93+ assertEq (thawEndTimestampAfter, 0 );
94+
95+ assertEq (tokenBalanceAfterSender, tokenBalanceBeforeSender + amountToWithdraw);
96+ assertEq (tokenBalanceAfterEscrow, tokenBalanceBeforeEscrow - amountToWithdraw);
97+ }
98+
7399 struct CollectPaymentData {
74100 uint256 escrowBalance;
75101 uint256 paymentsBalance;
@@ -105,7 +131,7 @@ contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
105131 }
106132
107133 vm.expectEmit (address (escrow));
108- emit IPaymentsEscrow.EscrowCollected (_payer, _collector, _receiver, _tokens);
134+ emit IPaymentsEscrow.EscrowCollected (_paymentType, _payer, _collector, _receiver, _tokens);
109135 escrow.collect (_paymentType, _payer, _receiver, _tokens, _dataService, _dataServiceCut);
110136
111137 // Calculate cuts
0 commit comments