Skip to content

Commit

Permalink
Fix tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@karim-en
karim-en committed May 29, 2024
1 parent 8c49070 commit 15cc4ba
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 22 deletions.
49 changes: 31 additions & 18 deletions token-locker/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -155,18 +155,18 @@ impl Contract {
#[init]
/// `prover_account`: NEAR account of the Near Prover contract;
/// `factory_address`: Ethereum address of the token factory contract, in hex.
pub fn new(prover_account: AccountId, factory_address: String) -> Self {
#[allow(deprecated)]
pub fn new(prover_account: AccountId, factory_address: String, whitelist_mode: bool) -> Self {
let mut contract = Self {
prover_account,
used_events: UnorderedSet::new(StorageKey::UsedEvents),
eth_factory_address: validate_eth_address(factory_address),
whitelist_tokens: UnorderedMap::new(StorageKey::WhitelistTokens),
whitelist_accounts: UnorderedSet::new(StorageKey::WhitelistAccounts),
is_whitelist_mode_enabled: true,
is_whitelist_mode_enabled: whitelist_mode,
};

contract.acl_init_super_admin(near_sdk::env::predecessor_account_id());
contract.acl_grant_role("DAO".to_owned(), near_sdk::env::predecessor_account_id());
contract
}

Expand Down Expand Up @@ -440,12 +440,18 @@ mod tests {
}
}

fn create_contract() -> Contract {
let whitelist_mode = true;
Contract::new(prover(), token_locker(), whitelist_mode)
}

#[test]
fn test_lock_unlock_token() {
set_env!(predecessor_account_id: accounts(0));
let mut contract = Contract::new(prover(), token_locker());
set_env!(predecessor_account_id: accounts(1));
let mut contract = create_contract();
contract.set_token_whitelist_mode(accounts(1), WhitelistMode::CheckToken);

set_env!(predecessor_account_id: accounts(1));
contract.ft_on_transfer(accounts(2), U128(1_000_000), ethereum_address_from_id(0));
contract.finish_deposit(
accounts(1).into(),
Expand All @@ -468,9 +474,10 @@ mod tests {
#[test]
fn test_lock_unlock_token_with_custom_recipient_message() {
set_env!(predecessor_account_id: accounts(0));
let mut contract = Contract::new(prover(), token_locker());
set_env!(predecessor_account_id: accounts(1));
let mut contract = create_contract();
contract.set_token_whitelist_mode(accounts(1), WhitelistMode::CheckToken);

set_env!(predecessor_account_id: accounts(1));
contract.ft_on_transfer(accounts(2), U128(1_000_000), ethereum_address_from_id(0));
contract.finish_deposit(
accounts(1).into(),
Expand Down Expand Up @@ -500,7 +507,7 @@ mod tests {
current_account_id: bridge_token_factory(),
predecessor_account_id: bridge_token_factory(),
);
let mut contract = Contract::new(prover(), token_locker());
let mut contract = create_contract();
assert!(contract
.acl_grant_role("PauseManager".to_owned(), pause_manager())
.unwrap());
Expand Down Expand Up @@ -532,33 +539,36 @@ mod tests {
#[should_panic(expected = "The token is blocked")]
fn test_blocked_token() {
set_env!(predecessor_account_id: accounts(0));
let mut contract = Contract::new(prover(), token_locker());
let mut contract = create_contract();

let token_account = accounts(1);
let sender_account = accounts(2);
set_env!(predecessor_account_id: token_account.clone());
contract.set_token_whitelist_mode(token_account, WhitelistMode::Blocked);
contract.set_token_whitelist_mode(token_account.clone(), WhitelistMode::Blocked);

set_env!(predecessor_account_id: token_account);
contract.ft_on_transfer(sender_account, U128(1_000_000), ethereum_address_from_id(0));
}

#[test]
#[should_panic(expected = "does not exist in the whitelist")]
fn test_account_not_in_whitelist() {
set_env!(predecessor_account_id: accounts(0));
let mut contract = Contract::new(prover(), token_locker());
let mut contract = create_contract();

let token_account = accounts(1);
let sender_account = accounts(2);
contract
.set_token_whitelist_mode(token_account.clone(), WhitelistMode::CheckAccountAndToken);

set_env!(predecessor_account_id: token_account);
contract.set_token_whitelist_mode(accounts(1), WhitelistMode::CheckAccountAndToken);
contract.ft_on_transfer(sender_account, U128(1_000_000), ethereum_address_from_id(0));
}

#[test]
#[should_panic(expected = "The token is not whitelisted")]
fn test_token_not_in_whitelist() {
set_env!(predecessor_account_id: accounts(0));
let mut contract = Contract::new(prover(), token_locker());
let mut contract = create_contract();

let token_account = accounts(1);
let sender_account = accounts(2);
Expand All @@ -569,7 +579,7 @@ mod tests {
#[test]
fn test_account_in_whitelist() {
set_env!(predecessor_account_id: accounts(0));
let mut contract = Contract::new(prover(), token_locker());
let mut contract = create_contract();

let token_account = accounts(1);
let sender_account = accounts(2);
Expand All @@ -585,7 +595,7 @@ mod tests {
#[should_panic(expected = "does not exist in the whitelist")]
fn test_remove_account_from_whitelist() {
set_env!(predecessor_account_id: accounts(0));
let mut contract = Contract::new(prover(), token_locker());
let mut contract = create_contract();

let token_account = accounts(1);
let sender_account = accounts(2);
Expand All @@ -600,7 +610,10 @@ mod tests {
ethereum_address_from_id(0),
);

set_env!(predecessor_account_id: accounts(0));
contract.remove_account_from_whitelist(token_account.clone(), sender_account.clone());

set_env!(predecessor_account_id: token_account.clone());
contract.ft_on_transfer(
sender_account.clone(),
U128(1_000_000),
Expand All @@ -611,7 +624,7 @@ mod tests {
#[test]
fn test_tokens_in_whitelist() {
set_env!(predecessor_account_id: accounts(0));
let mut contract = Contract::new(prover(), token_locker());
let mut contract = create_contract();

let whitelist_tokens = ["token1.near", "token2.near", "token3.near"];

Expand All @@ -630,7 +643,7 @@ mod tests {
#[test]
fn test_accounts_in_whitelist() {
set_env!(predecessor_account_id: accounts(0));
let mut contract = Contract::new(prover(), token_locker());
let mut contract = create_contract();

let whitelist_tokens = ["token1.near", "token2.near", "token3.near"];
let whitelist_accounts = ["account1.near", "account2.near", "account3.near"];
Expand Down
8 changes: 4 additions & 4 deletions token-locker/src/whitelist.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,12 @@ fn get_token_account_key(token: AccountId, account: AccountId) -> String {

#[near_bindgen]
impl Contract {
#[private]
#[access_control_any(roles(Role::DAO))]
pub fn set_token_whitelist_mode(&mut self, token: AccountId, mode: WhitelistMode) {
self.whitelist_tokens.insert(&token, &mode);
}

#[private]
#[access_control_any(roles(Role::DAO))]
pub fn add_account_to_whitelist(&mut self, token: AccountId, account: AccountId) {
assert!(
self.whitelist_tokens.get(&token).is_some(),
Expand All @@ -32,7 +32,7 @@ impl Contract {
.insert(&get_token_account_key(token, account));
}

#[private]
#[access_control_any(roles(Role::DAO))]
pub fn remove_account_from_whitelist(&mut self, token: AccountId, account: AccountId) -> bool {
self.whitelist_accounts
.remove(&get_token_account_key(token, account))
Expand Down Expand Up @@ -65,7 +65,7 @@ impl Contract {
}
}

#[private]
#[access_control_any(roles(Role::DAO))]
pub fn set_whitelist_mode_enabled(&mut self, enabled: bool) {
self.is_whitelist_mode_enabled = enabled;
}
Expand Down

0 comments on commit 15cc4ba

Please sign in to comment.