pub mod data;
pub mod error;
pub mod mocks;
pub mod state;
use data::account::UserAccountArgs;
use data::user::{UserData, UserDataArgs};
use error::UserStateError;
use ic_cdk::export::candid::Principal;
use state::UserState;
use std::cell::RefCell;
thread_local! {
static USER_STATE: RefCell<UserState> = RefCell::new(UserState::default());
}
pub fn initialize(wallet_canister_id: String) -> Result<(), UserStateError> {
let wallet_canister = Principal::from_text(wallet_canister_id).unwrap();
USER_STATE.with(|s| {
let mut mut_state = s.borrow_mut();
mut_state.init(wallet_canister);
Ok(())
})
}
pub fn get_wallet_canister() -> Result<Principal, UserStateError> {
USER_STATE.with(|s| {
let state = s.borrow();
Ok(state.wallet_canister)
})
}
pub fn get_owner() -> Result<Principal, UserStateError> {
USER_STATE.with(|s| {
let state = s.borrow();
Ok(state.owner)
})
}
pub fn change_owner(new_owner: Principal) -> Result<Principal, UserStateError> {
USER_STATE.with(|s| {
let mut mut_state = s.borrow_mut();
mut_state.change_owner(new_owner)
})
}
pub fn change_wallet_canister(wallet_canister: Principal) -> Result<Principal, UserStateError> {
USER_STATE.with(|s| {
let mut mut_state = s.borrow_mut();
mut_state.change_wallet_canister(wallet_canister)
})
}
pub fn add_user(
user: Principal,
user_args: UserDataArgs,
account_args: UserAccountArgs,
) -> Result<UserData, UserStateError> {
USER_STATE.with(|s| {
let mut mut_state = s.borrow_mut();
mut_state.create_user(user, user_args, account_args)
})
}
pub fn get_user(user: &Principal) -> Result<UserData, UserStateError> {
USER_STATE.with(|s| {
let state = s.borrow();
state.get_user(user).map(|user_data| user_data.clone())
})
}
pub fn with_user_state<T, F>(user: &Principal, callback: F) -> Result<T, UserStateError>
where
F: FnOnce(&UserData) -> T,
{
USER_STATE.with(|state| {
let state = state.borrow();
state.get_user(user).map(callback)
})
}
pub fn with_user_state_mut<T, F>(user: &Principal, callback: F) -> Result<T, UserStateError>
where
F: FnOnce(&mut UserData) -> T,
{
USER_STATE.with(|state| {
let mut state = state.borrow_mut();
state.get_user_mut(user).map(callback)
})
}
pub fn pre_upgrade() -> Result<(), candid::Error> {
USER_STATE.with(|s| {
let state = s.borrow();
ic_cdk::storage::stable_save((state.owner, state.wallet_canister, state.users.clone()))
})
}
pub fn post_upgrade() -> Result<(), candid::Error> {
USER_STATE.with(|s| {
let mut state = s.borrow_mut();
(state.owner, state.wallet_canister, state.users) =
ic_cdk::storage::stable_restore().unwrap();
Ok(())
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{mocks::*, state::UserState};
#[test]
fn test_initialize() {
let principal = random_principal();
assert!(initialize(principal.to_string()).is_ok());
USER_STATE.with(|s| {
let state = s.borrow();
assert_eq!(state.wallet_canister, principal);
});
}
#[test]
fn test_change_owner() {
let owner = random_principal();
assert!(change_owner(owner.clone()).is_ok());
USER_STATE.with(|s| {
let state = s.borrow();
assert_eq!(state.owner, owner);
});
}
#[test]
fn test_change_wallet_canister() {
owner_caller();
let wallet_canister = wallet_canister_principal();
assert!(change_wallet_canister(wallet_canister).is_ok());
USER_STATE.with(|s| {
let state = s.borrow();
assert_eq!(state.wallet_canister, wallet_canister);
});
}
#[test]
fn test_user_state_init() {
random_caller();
let mut user_state = UserState::default();
let principal = ic_caller();
random_caller();
user_state.init(principal);
assert_eq!(user_state.wallet_canister, principal);
let user = ic_caller();
user_state.init(principal);
random_caller();
assert!(user_state
.validate_caller_wallet_canister_or_user(&user)
.is_err());
}
#[test]
fn test_with_user_state() {
let principal = ic_caller();
initialize(principal.to_string()).unwrap();
let user_principal = random_principal();
let user_args = UserDataArgs {
balance: Some(100),
..UserDataArgs::default()
};
let account_args = UserAccountArgs {
name: Some("Account 1".to_owned()),
..UserAccountArgs::default()
};
set_caller(principal.clone());
let _ = add_user(
user_principal.clone(),
user_args.clone(),
account_args.clone(),
)
.unwrap();
set_caller(principal.clone());
let result = with_user_state(&user_principal, |user_data| user_data.balance).unwrap();
assert_eq!(result, user_args.balance.unwrap_or_default());
}
#[test]
fn test_with_user_state_mut() {
owner_caller();
let principal = wallet_canister_principal();
initialize(principal.to_string()).unwrap();
wallet_canister_caller();
let user_principal = random_principal();
let user_args = UserDataArgs {
balance: Some(100),
..UserDataArgs::default()
};
let account_args = UserAccountArgs {
name: Some("Account 1".to_owned()),
..UserAccountArgs::default()
};
let user_data = add_user(
user_principal.clone(),
user_args.clone(),
account_args.clone(),
)
.unwrap();
assert_eq!(user_data.balance, user_args.balance.unwrap_or_default());
let new_balance = 200;
let result = with_user_state_mut(&user_principal, |user_data| {
user_data.balance = new_balance;
user_data.balance
})
.unwrap();
assert_eq!(result, new_balance);
USER_STATE.with(|s| {
let state = s.borrow();
let stored_user_data = state.get_user(&user_principal).unwrap();
assert_eq!(stored_user_data.balance, new_balance);
});
}
#[test]
fn test_add_user() {
let principal = wallet_canister_principal();
owner_caller();
initialize(principal.to_string()).unwrap();
let user_principal = random_principal();
let user_args = UserDataArgs {
balance: Some(100),
..UserDataArgs::default()
};
let account_args = UserAccountArgs {
name: Some("Account 1".to_owned()),
..UserAccountArgs::default()
};
wallet_canister_caller();
let user_data = add_user(
user_principal.clone(),
user_args.clone(),
account_args.clone(),
)
.unwrap();
assert_eq!(user_data.balance, user_args.balance.unwrap_or_default());
assert_eq!(user_data.accounts.len(), 1);
assert_eq!(
user_data.accounts[0].name,
account_args.name.unwrap_or("Account 0".to_owned())
);
USER_STATE.with(|s| {
let state = s.borrow();
let stored_user_data = state.get_user(&user_principal).unwrap();
assert_eq!(
stored_user_data.balance,
user_args.balance.unwrap_or_default()
);
});
}
#[test]
fn test_get_user() {
let principal = wallet_canister_principal();
owner_caller();
initialize(principal.to_string()).unwrap();
let user_principal = random_principal();
let user_args = UserDataArgs {
balance: Some(100),
..UserDataArgs::default()
};
let account_args = UserAccountArgs {
name: Some("Account 1".to_owned()),
..UserAccountArgs::default()
};
wallet_canister_caller();
let _ = add_user(
user_principal.clone(),
user_args.clone(),
account_args.clone(),
)
.unwrap();
let user_data = get_user(&user_principal).unwrap();
assert_eq!(user_data.balance, user_args.balance.unwrap_or_default());
assert_eq!(user_data.accounts.len(), 1);
assert_eq!(
user_data.accounts[0].name,
account_args.name.unwrap_or("Account 0".to_owned())
);
USER_STATE.with(|s| {
let state = s.borrow();
let stored_user_data = state.get_user(&user_principal).unwrap();
assert_eq!(
stored_user_data.balance,
user_args.balance.unwrap_or_default()
);
});
}
}