diff --git a/rs/backend/src/accounts_store.rs b/rs/backend/src/accounts_store.rs index 39036f42a3..02e81e45fd 100644 --- a/rs/backend/src/accounts_store.rs +++ b/rs/backend/src/accounts_store.rs @@ -426,7 +426,13 @@ impl AccountsStore { // yet been stored, allowing us to set the principal (since originally we created accounts // without storing each user's principal). pub fn add_account(&mut self, caller: PrincipalId) -> bool { - self.assert_account_limit(); + self.add_account_with_limit(caller, ACCOUNT_LIMIT) + } + + /// Adds an account, up to the given account limit. + /// + /// The limit is a parameter so that a test can reach it without creating `ACCOUNT_LIMIT` accounts. + fn add_account_with_limit(&mut self, caller: PrincipalId, account_limit: u64) -> bool { let account_identifier = AccountIdentifier::from(caller); if let Some(account) = self.accounts_db.get(&account_identifier.to_vec()) { if account.principal.is_none() { @@ -437,6 +443,7 @@ impl AccountsStore { } false } else { + self.assert_account_limit(account_limit); let new_account = Account::new(caller, account_identifier); self.accounts_db.insert(account_identifier.to_vec(), new_account); @@ -881,11 +888,11 @@ impl AccountsStore { name.len() <= CANISTER_NAME_MAX_LENGTH } - fn assert_account_limit(&self) { + fn assert_account_limit(&self, account_limit: u64) { let db_accounts_len = self.accounts_db.len(); assert!( - db_accounts_len < ACCOUNT_LIMIT, - "Pre migration account limit exceeded {db_accounts_len}" + db_accounts_len < account_limit, + "Account limit exceeded {db_accounts_len}" ); } } diff --git a/rs/backend/src/accounts_store/tests.rs b/rs/backend/src/accounts_store/tests.rs index 05999a32ac..ff79f573f7 100644 --- a/rs/backend/src/accounts_store/tests.rs +++ b/rs/backend/src/accounts_store/tests.rs @@ -89,6 +89,63 @@ fn create_sub_account_limit_exceeded() { assert_eq!(MAX_SUB_ACCOUNTS_PER_ACCOUNT, sub_accounts.len()); } +#[test] +#[should_panic(expected = "Account limit exceeded 2")] +fn add_account_at_limit_traps_for_new_account() { + let principal = PrincipalId::from_str(TEST_ICRC1_ACCOUNT_3).unwrap(); + let mut store = setup_test_store(); + + assert_eq!(2, store.accounts_db.len()); + + store.add_account_with_limit(principal, 2); +} + +#[test] +fn add_account_at_limit_returns_false_for_existing_account() { + let principal = PrincipalId::from_str(TEST_ICRC1_ACCOUNT_1).unwrap(); + let mut store = setup_test_store(); + + assert_eq!(2, store.accounts_db.len()); + + assert!(!store.add_account_with_limit(principal, 2)); + + assert_eq!(2, store.accounts_db.len()); + assert!(store.get_account(principal).is_some()); +} + +#[test] +fn add_account_at_limit_backfills_legacy_principal() { + let principal = PrincipalId::from_str(TEST_ICRC1_ACCOUNT_3).unwrap(); + let account_identifier = AccountIdentifier::from(principal); + let mut store = setup_test_store(); + + // An account created before the principal was stored. + let mut account = Account::new(principal, account_identifier); + account.principal = None; + store.accounts_db.insert(account_identifier.to_vec(), account); + + assert_eq!(3, store.accounts_db.len()); + assert!(store.get_account(principal).is_none()); + + assert!(!store.add_account_with_limit(principal, 3)); + + assert_eq!(3, store.accounts_db.len()); + assert_eq!(principal, store.get_account(principal).unwrap().principal); +} + +#[test] +fn add_account_below_limit_creates_account() { + let principal = PrincipalId::from_str(TEST_ICRC1_ACCOUNT_3).unwrap(); + let mut store = setup_test_store(); + + assert_eq!(2, store.accounts_db.len()); + + assert!(store.add_account_with_limit(principal, 3)); + + assert_eq!(3, store.accounts_db.len()); + assert!(store.get_account(principal).is_some()); +} + #[test] fn create_sub_account_account_not_found() { let principal = PrincipalId::from_str(TEST_ICRC1_ACCOUNT_3).unwrap();