Add tests to cover upgrading to Diamond by paying in CHAT

backend/integration_tests/src/client/mod.rs

@@ -32,12 +32,8 @@ pub fn create_canister(env: &mut PocketIc, controller: Principal) -> CanisterId
 }
 
 pub fn create_canister_with_id(env: &mut PocketIc, controller: Principal, canister_id: &str) -> CanisterId {
-    let canister_id = env
-        .create_canister_with_id(
-            Some(controller),
-            None,
-            Principal::from_text(canister_id).expect("Invalid canister ID"),
-        )
+    let canister_id = canister_id.try_into().expect("Invalid canister ID");
+    env.create_canister_with_id(Some(controller), None, canister_id)
         .expect("Create canister with ID failed");
     env.add_cycles(canister_id, INIT_CYCLES_BALANCE);
     canister_id
@@ -94,25 +90,7 @@ pub fn execute_update_no_response<P: CandidType>(
 
 pub fn register_diamond_user(env: &mut PocketIc, canister_ids: &CanisterIds, controller: Principal) -> User {
     let user = local_user_index::happy_path::register_user(env, canister_ids.local_user_index);
-
-    icrc1::happy_path::transfer(
-        env,
-        controller,
-        canister_ids.icp_ledger,
-        user.user_id.into(),
-        10_000_000_000u64,
-    );
-
-    user_index::happy_path::pay_for_diamond_membership(
-        env,
-        user.principal,
-        canister_ids.user_index,
-        DiamondMembershipPlanDuration::OneMonth,
-        true,
-    );
-
-    tick_many(env, 3);
-
+    upgrade_user(&user, env, canister_ids, controller);
     user
 }
 
@@ -130,6 +108,7 @@ pub fn upgrade_user(user: &User, env: &mut PocketIc, canister_ids: &CanisterIds,
         user.principal,
         canister_ids.user_index,
         DiamondMembershipPlanDuration::OneMonth,
+        false,
         true,
     );
 

backend/integration_tests/src/client/user_index.rs

@@ -77,6 +77,7 @@ pub mod happy_path {
         sender: Principal,
         canister_id: CanisterId,
         duration: DiamondMembershipPlanDuration,
+        pay_in_chat: bool,
         recurring: bool,
     ) -> DiamondMembershipDetails {
         let response = super::pay_for_diamond_membership(
@@ -85,8 +86,8 @@ pub mod happy_path {
             canister_id,
             &user_index_canister::pay_for_diamond_membership::Args {
                 duration,
-                token: Cryptocurrency::InternetComputer,
-                expected_price_e8s: duration.icp_price_e8s(),
+                token: if pay_in_chat { Cryptocurrency::CHAT } else { Cryptocurrency::InternetComputer },
+                expected_price_e8s: if pay_in_chat { duration.chat_price_e8s() } else { duration.icp_price_e8s() },
                 recurring,
             },
         );

backend/integration_tests/src/diamond_membership_tests.rs

@@ -10,10 +10,12 @@ use types::{Cryptocurrency, DiamondMembershipPlanDuration, DiamondMembershipSubs
 use utils::consts::SNS_GOVERNANCE_CANISTER_ID;
 use utils::time::MINUTE_IN_MS;
 
-#[test_case(false)]
-#[test_case(true)]
+#[test_case(true, false)]
+#[test_case(true, true)]
+#[test_case(false, false)]
+#[test_case(false, true)]
 #[serial]
-fn can_upgrade_to_diamond(lifetime: bool) {
+fn can_upgrade_to_diamond(pay_in_chat: bool, lifetime: bool) {
     let mut wrapper = ENV.deref().get();
     let TestEnv {
         env,
@@ -22,17 +24,13 @@ fn can_upgrade_to_diamond(lifetime: bool) {
         ..
     } = wrapper.env();
 
-    let init_treasury_balance = client::icrc1::happy_path::balance_of(env, canister_ids.icp_ledger, SNS_GOVERNANCE_CANISTER_ID);
+    let ledger = if pay_in_chat { canister_ids.chat_ledger } else { canister_ids.icp_ledger };
+
+    let init_treasury_balance = client::icrc1::happy_path::balance_of(env, ledger, SNS_GOVERNANCE_CANISTER_ID);
 
     let user = client::local_user_index::happy_path::register_user(env, canister_ids.local_user_index);
 
-    client::icrc1::happy_path::transfer(
-        env,
-        *controller,
-        canister_ids.icp_ledger,
-        user.user_id.into(),
-        1_000_000_000u64,
-    );
+    client::icrc1::happy_path::transfer(env, *controller, ledger, user.user_id.into(), 10_000_000_000u64);
 
     let now = now_millis(env);
 
@@ -49,6 +47,7 @@ fn can_upgrade_to_diamond(lifetime: bool) {
         user.principal,
         canister_ids.user_index,
         duration,
+        pay_in_chat,
         false,
     );
 
@@ -69,14 +68,20 @@ fn can_upgrade_to_diamond(lifetime: bool) {
         .subscription
         .is_active());
 
-    let new_balance = client::icrc1::happy_path::balance_of(env, canister_ids.icp_ledger, user.user_id.into());
-    assert_eq!(new_balance, 1_000_000_000 - duration.icp_price_e8s());
+    let (expected_price, transfer_fee) = if pay_in_chat {
+        (duration.chat_price_e8s(), Cryptocurrency::CHAT.fee().unwrap())
+    } else {
+        (duration.icp_price_e8s(), Cryptocurrency::InternetComputer.fee().unwrap())
+    };
 
-    let treasury_balance = client::icrc1::happy_path::balance_of(env, canister_ids.icp_ledger, SNS_GOVERNANCE_CANISTER_ID);
+    let new_balance = client::icrc1::happy_path::balance_of(env, ledger, user.user_id.into());
+    assert_eq!(new_balance, 10_000_000_000 - expected_price);
+
+    let treasury_balance = client::icrc1::happy_path::balance_of(env, ledger, SNS_GOVERNANCE_CANISTER_ID);
 
     assert_eq!(
         treasury_balance - init_treasury_balance,
-        duration.icp_price_e8s() - (2 * Cryptocurrency::InternetComputer.fee().unwrap()) as u64
+        expected_price - (2 * transfer_fee) as u64
     );
 }
 

backend/integration_tests/src/lib.rs

@@ -78,6 +78,7 @@ pub struct CanisterIds {
     pub cycles_dispenser: CanisterId,
     pub registry: CanisterId,
     pub icp_ledger: CanisterId,
+    pub chat_ledger: CanisterId,
     pub cycles_minting_canister: CanisterId,
 }
 

backend/integration_tests/src/registry_tests.rs

@@ -1,6 +1,6 @@
 use crate::env::ENV;
 use crate::rng::{random_principal, random_string};
-use crate::setup::install_icrc1_ledger;
+use crate::setup::install_icrc_ledger;
 use crate::utils::now_millis;
 use crate::{client, TestEnv};
 use registry_canister::TokenStandard;
@@ -17,12 +17,13 @@ fn add_token_succeeds() {
         ..
     } = wrapper.env();
 
-    let ledger_canister_id = install_icrc1_ledger(
+    let ledger_canister_id = install_icrc_ledger(
         env,
         *controller,
         "ABC Token".to_string(),
         "ABC".to_string(),
         10_000,
+        None,
         Vec::new(),
     );
 
@@ -105,12 +106,13 @@ fn update_token_succeeds() {
         ..
     } = wrapper.env();
 
-    let ledger_canister_id = install_icrc1_ledger(
+    let ledger_canister_id = install_icrc_ledger(
         env,
         *controller,
         "ABC Token".to_string(),
         "ABC".to_string(),
         10_000,
+        None,
         Vec::new(),
     );
 

backend/integration_tests/src/setup.rs

@@ -38,7 +38,11 @@ pub fn setup_new_env() -> TestEnv {
         ", &path, &env::current_dir().map(|x| x.display().to_string()).unwrap_or_else(|_| "an unknown directory".to_string()));
     }
 
-    let mut env = PocketIcBuilder::new().with_nns_subnet().with_application_subnet().build();
+    let mut env = PocketIcBuilder::new()
+        .with_nns_subnet()
+        .with_sns_subnet()
+        .with_application_subnet()
+        .build();
     let controller = random_principal();
     let canister_ids = install_canisters(&mut env, controller);
 
@@ -56,6 +60,15 @@ fn install_canisters(env: &mut PocketIc, controller: Principal) -> CanisterIds {
     let cycles_minting_canister_id = create_canister_with_id(env, controller, "rkp4c-7iaaa-aaaaa-aaaca-cai");
     let sns_wasm_canister_id = create_canister_with_id(env, controller, "qaa6y-5yaaa-aaaaa-aaafa-cai");
     let nns_index_canister_id = create_canister_with_id(env, controller, "qhbym-qaaaa-aaaaa-aaafq-cai");
+    let chat_ledger_canister_id = install_icrc_ledger(
+        env,
+        controller,
+        "OpenChat".to_string(),
+        "CHAT".to_string(),
+        100000,
+        Some("2ouva-viaaa-aaaaq-aaamq-cai"),
+        Vec::new(),
+    );
 
     let user_index_canister_id = create_canister(env, controller);
     let group_index_canister_id = create_canister(env, controller);
@@ -364,16 +377,18 @@ fn install_canisters(env: &mut PocketIc, controller: Principal) -> CanisterIds {
         cycles_dispenser: cycles_dispenser_canister_id,
         registry: registry_canister_id,
         icp_ledger: nns_ledger_canister_id,
+        chat_ledger: chat_ledger_canister_id,
         cycles_minting_canister: cycles_minting_canister_id,
     }
 }
 
-pub fn install_icrc1_ledger(
+pub fn install_icrc_ledger(
     env: &mut PocketIc,
     controller: Principal,
     token_name: String,
     token_symbol: String,
     transfer_fee: u64,
+    canister_id: Option<&str>,
     initial_balances: Vec<(Account, u64)>,
 ) -> CanisterId {
     #[derive(CandidType)]
@@ -413,7 +428,11 @@ pub fn install_icrc1_ledger(
         },
     });
 
-    let canister_id = create_canister(env, controller);
+    let canister_id = if let Some(id) = canister_id {
+        create_canister_with_id(env, controller, id)
+    } else {
+        create_canister(env, controller)
+    };
     install_canister(env, controller, canister_id, wasms::ICRC_LEDGER.clone(), args);
 
     canister_id