ic-testkit 0.7.6

PocketIC-oriented test utilities for IC canister tests
Documentation
use candid::Principal;
use ic_testkit::pic::{
    CanisterSnapshotTarget, ControllerSnapshotError, PocketIc, PocketIcCapturedSnapshotExt,
    PocketIcSnapshotExt, SnapshotRestoreFunding,
};
use pocket_ic::CanisterSettings;

#[test]
fn explicit_snapshot_senders_support_mixed_controller_sets_without_fallback() {
    let pocket_ic = PocketIc::new();
    let explicit_controller = Principal::from_slice(&[42]);
    let anonymous_canister = pocket_ic.create_canister();
    let explicit_canister = pocket_ic.create_canister_with_settings(
        None,
        Some(CanisterSettings {
            controllers: Some(vec![explicit_controller]),
            ..CanisterSettings::default()
        }),
    );
    for (canister_id, sender) in [
        (anonymous_canister, None),
        (explicit_canister, Some(explicit_controller)),
    ] {
        pocket_ic.install_canister(canister_id, b"\0asm\x01\0\0\0".to_vec(), vec![], sender);
        pocket_ic
            .stop_canister(canister_id, sender)
            .expect("stop mixed-controller canister before snapshot capture");
    }

    let snapshots = pocket_ic
        .capture_snapshots_with_senders([
            CanisterSnapshotTarget::new(anonymous_canister, None),
            CanisterSnapshotTarget::new(explicit_canister, Some(explicit_controller)),
        ])
        .expect("capture mixed-controller snapshots with exact senders");

    assert_eq!(snapshots.len(), 2);
    let mut expected = vec![anonymous_canister, explicit_canister];
    expected.sort();
    assert_eq!(snapshots.canister_ids().collect::<Vec<_>>(), expected);
    pocket_ic
        .restore_snapshots_with_captured_senders(&snapshots)
        .expect("restore mixed-controller snapshots with exact captured senders");
}

#[test]
fn failed_snapshot_set_capture_cleans_up_earlier_snapshots() {
    let pocket_ic = PocketIc::new();
    let captured_first = pocket_ic.create_canister();
    pocket_ic.install_canister(captured_first, b"\0asm\x01\0\0\0".to_vec(), vec![], None);
    pocket_ic
        .stop_canister(captured_first, None)
        .expect("stop canister before snapshot capture");

    let missing_canister = pocket_ic.create_canister();
    pocket_ic
        .stop_canister(missing_canister, None)
        .expect("stop canister before deletion");
    pocket_ic
        .delete_canister(missing_canister, None)
        .expect("delete canister used for capture failure");
    assert!(captured_first < missing_canister);

    let error = pocket_ic
        .capture_controller_snapshots(Principal::anonymous(), [captured_first, missing_canister])
        .expect_err("missing canister should fail the snapshot set");

    match error {
        ControllerSnapshotError::CaptureFailed {
            canister_id,
            attempts,
            cleanup_failures,
        } => {
            assert_eq!(canister_id, missing_canister, "attempts: {attempts:?}");
            assert!(cleanup_failures.is_empty());
        }
        other => panic!("unexpected snapshot error: {other}"),
    }

    let remaining = pocket_ic
        .list_canister_snapshots(captured_first, None)
        .expect("list snapshots after rollback");
    assert!(remaining.is_empty());
}

#[test]
fn snapshot_restore_only_adds_cycles_when_explicitly_requested() {
    let pocket_ic = PocketIc::new();
    let canister_id = pocket_ic.create_canister();
    pocket_ic.install_canister(canister_id, b"\0asm\x01\0\0\0".to_vec(), vec![], None);
    pocket_ic
        .stop_canister(canister_id, None)
        .expect("stop canister before snapshot capture");

    let snapshots = pocket_ic
        .capture_controller_snapshots(Principal::anonymous(), [canister_id])
        .expect("capture snapshot set");
    let balance_before_restore = pocket_ic.cycle_balance(canister_id);

    pocket_ic
        .restore_controller_snapshots(Principal::anonymous(), &snapshots)
        .expect("restore without funding");
    let balance_after_preserved_restore = pocket_ic.cycle_balance(canister_id);
    assert!(
        balance_after_preserved_restore <= balance_before_restore,
        "default restore must not add cycles"
    );

    let minimum_cycles = balance_after_preserved_restore + 10_000_000_000_000;
    pocket_ic
        .restore_controller_snapshots_with_funding(
            Principal::anonymous(),
            &snapshots,
            SnapshotRestoreFunding::TopUpTo { minimum_cycles },
        )
        .expect("restore with explicit funding");
    assert!(
        pocket_ic.cycle_balance(canister_id) > balance_after_preserved_restore,
        "explicit top-up should increase the post-restore balance"
    );
}