@@ -174,4 +174,82 @@ contract EntropyAuthorized is Test, EntropyTestUtils {
174
174
vm.expectRevert (EntropyErrors.Unauthorized.selector );
175
175
random.acceptAdmin ();
176
176
}
177
+
178
+ // Helper function to setup contract with fees
179
+ function setupContractWithFees (
180
+ uint128 feeAmount ,
181
+ uint numRequests
182
+ ) internal returns (uint128 totalFees ) {
183
+ // Register provider1
184
+ bytes32 [] memory hashChain = generateHashChain (provider1, 0 , 100 );
185
+ vm.prank (provider1);
186
+ random.register (0 , hashChain[0 ], hex "0100 " , 100 , "" );
187
+
188
+ // Set Pyth fee
189
+ vm.prank (admin);
190
+ random.setPythFee (feeAmount);
191
+
192
+ // Make requests to accrue fees
193
+ bytes32 userCommitment = random.constructUserCommitment (
194
+ bytes32 (uint256 (42 ))
195
+ );
196
+ vm.deal (address (this ), feeAmount * numRequests);
197
+ for (uint i = 0 ; i < numRequests; i++ ) {
198
+ random.request {value: feeAmount}(provider1, userCommitment, false );
199
+ }
200
+
201
+ totalFees = uint128 (feeAmount * numRequests);
202
+ assertEq (random.getAccruedPythFees (), totalFees);
203
+ return totalFees;
204
+ }
205
+
206
+ function testWithdrawFeeByAdmin () public {
207
+ uint128 totalFees = setupContractWithFees (10 , 5 );
208
+
209
+ address targetAddress = address (123 );
210
+ uint128 withdrawAmount = 30 ;
211
+
212
+ vm.prank (admin);
213
+ random.withdrawFee (targetAddress, withdrawAmount);
214
+
215
+ assertEq (random.getAccruedPythFees (), totalFees - withdrawAmount);
216
+ assertEq (targetAddress.balance, withdrawAmount);
217
+ }
218
+
219
+ function testWithdrawFeeByOwner () public {
220
+ uint128 totalFees = setupContractWithFees (10 , 5 );
221
+
222
+ address targetAddress = address (123 );
223
+ uint128 withdrawAmount = 30 ;
224
+
225
+ vm.prank (owner);
226
+ random.withdrawFee (targetAddress, withdrawAmount);
227
+
228
+ assertEq (random.getAccruedPythFees (), totalFees - withdrawAmount);
229
+ assertEq (targetAddress.balance, withdrawAmount);
230
+ }
231
+
232
+ function testWithdrawFeeByUnauthorized () public {
233
+ setupContractWithFees (10 , 5 );
234
+
235
+ vm.prank (admin2);
236
+ vm.expectRevert (EntropyErrors.Unauthorized.selector );
237
+ random.withdrawFee (address (123 ), 30 );
238
+ }
239
+
240
+ function testWithdrawFeeInsufficientBalance () public {
241
+ uint128 totalFees = setupContractWithFees (10 , 5 );
242
+
243
+ vm.prank (admin);
244
+ vm.expectRevert (EntropyErrors.InsufficientFee.selector );
245
+ random.withdrawFee (address (123 ), totalFees + 10 );
246
+ }
247
+
248
+ function testWithdrawFeeToZeroAddress () public {
249
+ setupContractWithFees (10 , 5 );
250
+
251
+ vm.prank (admin);
252
+ vm.expectRevert ("targetAddress is zero address " );
253
+ random.withdrawFee (address (0 ), 30 );
254
+ }
177
255
}
0 commit comments