From b58c09a203cfacf0263be63dea6f0e078d664535 Mon Sep 17 00:00:00 2001
From: Max Murphy <max@winning.black>
Date: Mon, 9 Sep 2024 15:30:07 +0200
Subject: [PATCH] ++

---
 src/example/paid_service/src/lib.rs        |  6 +++---
 src/example/paid_service/tests/it/icrc2.rs | 22 +++++++++++++++++-----
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/src/example/paid_service/src/lib.rs b/src/example/paid_service/src/lib.rs
index 9ccda71..0e938e8 100644
--- a/src/example/paid_service/src/lib.rs
+++ b/src/example/paid_service/src/lib.rs
@@ -28,9 +28,9 @@ async fn cost_1000_attached_cycles() -> Result<String, PaymentError> {
     Ok("Yes, you paid 1000 cycles!".to_string())
 }
 
-/// An API method that requires cycles to be provided by the user using an ICRC-2 approve.
+/// An API method that requires 1 billion cycles to be provided by the user using an ICRC-2 approve.
 #[update()]
-async fn cost_1000_icrc2_from_caller() -> Result<String, PaymentError> {
+async fn cost_1b_icrc2_from_caller() -> Result<String, PaymentError> {
     let guard = Icrc2FromCaller {
         own_canister_id: own_canister_id(),
         payer: Account {
@@ -39,7 +39,7 @@ async fn cost_1000_icrc2_from_caller() -> Result<String, PaymentError> {
         },
         ledger_canister_id: payment_ledger(),
     };
-    guard.deduct(1000).await?;
+    guard.deduct(1_000_000_000).await?;
     Ok("Yes, you paid 1000 cycles!".to_string())
 }
 
diff --git a/src/example/paid_service/tests/it/icrc2.rs b/src/example/paid_service/tests/it/icrc2.rs
index d652285..752c522 100644
--- a/src/example/paid_service/tests/it/icrc2.rs
+++ b/src/example/paid_service/tests/it/icrc2.rs
@@ -138,16 +138,17 @@ fn icrc2_payment_works() {
     );
     // .. Get enough to play with lots of transactions.
     const LEDGER_FEE: u128 = 100_000_000; // The documented fee: https://internetcomputer.org/docs/current/developer-docs/defi/cycles/cycles-ledger#fees
-    let mut remainder = 100u128; // Multiple of fees we have left to play with.
-    setup.fund_user(LEDGER_FEE * remainder);
+    let mut expected_user_balance = 100_000_000_000; // Lots of funds to play with.
+    setup.fund_user(expected_user_balance);
     setup.assert_user_balance_eq(
-        LEDGER_FEE * remainder,
+        expected_user_balance,
         "Test setup failed when providing the user with funds".to_string(),
     );
     // Exercise the protocol...
-    let api_method = "cost_1000_icrc2_from_caller";
-    let api_fee = 1_000u128;
+    let api_method = "cost_1b_icrc2_from_caller";
+    let api_fee = 1_000_000_000u128;
     for payment in (api_fee - 5)..(api_fee + 5) {
+        // Pre-approve payment
         setup
             .ledger
             .icrc_2_approve(
@@ -163,6 +164,10 @@ fn icrc2_payment_works() {
             )
             .expect("Failed to call the ledger to approve")
             .expect("Failed to approve the paid service to spend the user's ICRC-2 tokens");
+        // Check the balance beforehand
+        let service_canister_cycles_before =
+            setup.pic.cycle_balance(setup.paid_service.canister_id);
+        // Call the API
         let response: Result<String, PaymentError> = setup
             .paid_service
             .update(setup.user, api_method, ())
@@ -186,6 +191,13 @@ fn icrc2_payment_works() {
                 "Should have succeeded with {} cycles attached",
                 payment
             );
+            let service_canister_cycles_after =
+                setup.pic.cycle_balance(setup.paid_service.canister_id);
+            assert!(
+                service_canister_cycles_after > service_canister_cycles_before,
+                "The service canister needs to charge more to cover its cycle cost!  Loss: {}",
+                service_canister_cycles_before - service_canister_cycles_after
+            );
         }
     }
 }