@@ -1361,6 +1361,153 @@ async fn test_pending_payments_block_usage() {
13611361 . await ;
13621362}
13631363
1364+ #[ tokio:: test]
1365+ async fn test_payment_tracker_usage_tracking ( ) {
1366+ // This test verifies that:
1367+ // 1. Pricing increases with concurrency (usage percentage)
1368+ // 2. After requests complete, usage tracking correctly returns to 0%
1369+ // (verifying that PaymentUsageGuard properly deregisters when requests finish)
1370+
1371+ crate :: common:: setup_logging ( ) ;
1372+ let ( testnet, _validator_collection, actions, node_set) = setup_testnet_for_payments ( ) . await ;
1373+ let realm_id = ethers:: types:: U256 :: from ( 1 ) ;
1374+ let node_set = get_identity_pubkeys_from_node_set ( & node_set) . await ;
1375+
1376+ let self_pay_user = EndUser :: new ( & testnet) ;
1377+ self_pay_user
1378+ . set_wallet_balance ( INITIAL_FUNDING_AMOUNT )
1379+ . await ;
1380+
1381+ // Get initial price at 0% usage
1382+ // 3 is the SignSessionKey product ID
1383+ let product_id = 1 ;
1384+ let initial_price = self_pay_user. first_node_price_from_feed ( product_id) . await ;
1385+ info ! ( "Initial price at 0% usage: {}" , initial_price) ;
1386+
1387+ // Prepare valid encryption parameters for successful requests
1388+ let test_encryption_parameters = prepare_test_encryption_parameters ( ) ;
1389+
1390+ let resource_ability_requests = vec ! [ LitResourceAbilityRequest {
1391+ resource: LitResourceAbilityRequestResource {
1392+ resource: format!(
1393+ "{}/{}" ,
1394+ test_encryption_parameters. hashed_access_control_conditions,
1395+ test_encryption_parameters. data_to_encrypt_hash
1396+ ) ,
1397+ resource_prefix: LitResourcePrefix :: ACC . to_string( ) ,
1398+ } ,
1399+ ability: LitAbility :: AccessControlConditionDecryption . to_string( ) ,
1400+ } ] ;
1401+
1402+ // Fund the user with enough balance for multiple requests
1403+ // Use a multiplier for funding to ensure we have enough for concurrent requests
1404+ let funding_amount = initial_price * 1000 * NUM_STAKED_VALIDATORS ;
1405+ self_pay_user. deposit_to_wallet_ledger ( funding_amount) . await ;
1406+
1407+ // Step 1: Make concurrent requests to increase usage and verify pricing increases
1408+ info ! ( "Step 1: Making concurrent requests to increase usage" ) ;
1409+ let num_concurrent_requests = 30 ;
1410+ let mut handles = Vec :: new ( ) ;
1411+
1412+ for i in 0 ..num_concurrent_requests {
1413+ let session_sigs_and_node_set = get_session_sigs_for_auth (
1414+ & node_set,
1415+ resource_ability_requests. clone ( ) ,
1416+ Some ( self_pay_user. wallet . clone ( ) ) ,
1417+ None ,
1418+ Some ( initial_price) ,
1419+ ) ;
1420+
1421+ let params = test_encryption_parameters. clone ( ) ;
1422+ let epoch = actions. get_current_epoch ( realm_id) . await . as_u64 ( ) ;
1423+
1424+ let handle = tokio:: spawn ( async move {
1425+ // Add small delay to ensure requests are truly concurrent
1426+ tokio:: time:: sleep ( tokio:: time:: Duration :: from_millis ( i * 10 ) ) . await ;
1427+ retrieve_decryption_key_session_sigs (
1428+ params,
1429+ & session_sigs_and_node_set,
1430+ epoch,
1431+ DEFAULT_KEY_SET_NAME ,
1432+ )
1433+ . await
1434+ } ) ;
1435+ handles. push ( handle) ;
1436+ }
1437+
1438+ // Wait a bit for requests to register usage
1439+ tokio:: time:: sleep ( tokio:: time:: Duration :: from_millis ( 200 ) ) . await ;
1440+
1441+ // Check that price has increased due to concurrency
1442+ let price_with_concurrency = self_pay_user. first_node_price_from_feed ( product_id) . await ;
1443+ info ! (
1444+ "Price with {} concurrent requests: {}" ,
1445+ num_concurrent_requests, price_with_concurrency
1446+ ) ;
1447+
1448+ // Price should be higher than initial
1449+ assert ! (
1450+ price_with_concurrency >= initial_price,
1451+ "Price should increase or stay the same with concurrent requests. Initial: {}, With concurrency: {}" ,
1452+ initial_price,
1453+ price_with_concurrency
1454+ ) ;
1455+
1456+ // Wait for all requests to complete
1457+ info ! ( "Waiting for all concurrent requests to complete" ) ;
1458+ for handle in handles {
1459+ let _ = handle. await ;
1460+ }
1461+
1462+ // Step 2: Wait for all requests to fully complete and verify usage returns to 0%
1463+ info ! ( "Step 2: Waiting for usage to return to 0%" ) ;
1464+
1465+ // Give some time for all guards to drop and deregister
1466+ tokio:: time:: sleep ( tokio:: time:: Duration :: from_secs ( 2 ) ) . await ;
1467+
1468+ // Check that price has returned to initial (or close to it)
1469+ // Note: There might be some delay in price feed updates, so we check multiple times
1470+ let mut final_price = self_pay_user. first_node_price_from_feed ( product_id) . await ;
1471+ let mut attempts = 0 ;
1472+ const MAX_ATTEMPTS : u32 = 10 ;
1473+
1474+ while final_price > initial_price && attempts < MAX_ATTEMPTS {
1475+ tokio:: time:: sleep ( tokio:: time:: Duration :: from_millis ( 500 ) ) . await ;
1476+ final_price = self_pay_user. first_node_price_from_feed ( product_id) . await ;
1477+ attempts += 1 ;
1478+ info ! (
1479+ "Attempt {}: Final price: {}, Initial price: {}" ,
1480+ attempts, final_price, initial_price
1481+ ) ;
1482+ }
1483+
1484+ info ! ( "Final price after all requests: {}" , final_price) ;
1485+ info ! ( "Initial price: {}" , initial_price) ;
1486+
1487+ // The final price should be close to the initial price (within reasonable tolerance)
1488+ // This verifies that usage tracking correctly returned to 0% after all requests completed
1489+ // Note: We allow some tolerance because price feed updates might have slight delays
1490+ let price_difference = if final_price > initial_price {
1491+ final_price - initial_price
1492+ } else {
1493+ U256 :: zero ( )
1494+ } ;
1495+
1496+ // Allow up to 5% difference to account for price feed update delays
1497+ let tolerance = initial_price / 20 ;
1498+
1499+ assert ! (
1500+ price_difference <= tolerance,
1501+ "Price should return close to initial after all requests complete (including exceptions). \
1502+ \
1503+ ,
1504+ initial_price,
1505+ final_price,
1506+ price_difference,
1507+ tolerance
1508+ ) ;
1509+ }
1510+
13641511async fn setup_testnet_for_payments ( ) -> ( Testnet , ValidatorCollection , Actions , Vec < NodeSet > ) {
13651512 do_setup_testnet_for_payments ( true ) . await
13661513}
0 commit comments