chore: storage management tests

Cargo.toml

@@ -13,12 +13,16 @@ repository = "https://github.com/NEARFoundation/near-sdk-contract-tools"
 version = "1.1.0"
 
 [workspace.dependencies]
+# normal dependencies
 near-sdk = { version = "4.1.1", default-features = false }
 near-sdk-contract-tools-macros = { version = "=1.1.0", path = "./macros" }
 serde = "1.0.144"
 serde_json = "1.0.85"
 thiserror = "1.0.35"
 
+# dev-dependencies
+pretty_assertions = "1.4.0"
+
 [dependencies]
 near-sdk.workspace = true
 near-sdk-contract-tools-macros.workspace = true

workspaces-tests-utils/Cargo.toml

@@ -7,3 +7,4 @@ publish = false
 [target.'cfg(not(windows))'.dependencies]
 near-sdk.workspace = true
 near-workspaces = "0.8"
+pretty_assertions.workspace = true

workspaces-tests-utils/src/lib.rs

@@ -1,8 +1,9 @@
 #![allow(missing_docs)]
 #![cfg(not(windows))]
 
-use near_sdk::{serde::de::DeserializeOwned, serde_json::json};
-use near_workspaces::{result::ExecutionFinalResult, Account, Contract};
+use near_sdk::{json_types::U128, serde::de::DeserializeOwned, serde_json::json};
+use near_workspaces::{result::ExecutionFinalResult, Account, AccountId, Contract};
+use pretty_assertions::assert_eq;
 
 pub async fn nft_token<T: DeserializeOwned>(contract: &Contract, token_id: &str) -> Option<T> {
     contract
@@ -14,6 +15,17 @@ pub async fn nft_token<T: DeserializeOwned>(contract: &Contract, token_id: &str)
         .unwrap()
 }
 
+pub async fn ft_balance_of(contract: &Contract, account: &AccountId) -> u128 {
+    contract
+        .view("ft_balance_of")
+        .args_json(json!({ "account_id": account }))
+        .await
+        .unwrap()
+        .json::<U128>()
+        .map(u128::from)
+        .unwrap()
+}
+
 pub struct Setup {
     pub contract: Contract,
     pub accounts: Vec<Account>,

workspaces-tests/Cargo.toml

@@ -72,7 +72,7 @@ thiserror.workspace = true
 [dev-dependencies]
 near-crypto = "0.15.0"
 tokio = "1.21.1"
-pretty_assertions = "1.4.0"
+pretty_assertions.workspace = true
 
 [target.'cfg(not(windows))'.dev-dependencies]
 near-workspaces = "0.8"

workspaces-tests/src/bin/fungible_token_nep145.rs

@@ -3,52 +3,76 @@
 // Ignore
 pub fn main() {}
 
+use std::cmp::Ordering;
+
 use near_sdk::{
     borsh::{self, BorshDeserialize, BorshSerialize},
     env,
     json_types::U128,
-    near_bindgen, PanicOnDefault,
-};
-use near_sdk_contract_tools::{
-    standard::{nep141::*, nep145::*},
-    FungibleToken, Nep145,
+    near_bindgen, require, AccountId, PanicOnDefault,
 };
+use near_sdk_contract_tools::{ft::*, standard::nep145::*, Nep145};
 
 #[derive(PanicOnDefault, BorshSerialize, BorshDeserialize, FungibleToken, Nep145)]
-#[fungible_token(
-    name = "My Fungible Token w/ Storage Management",
-    symbol = "MYFT",
-    decimals = 24
-)]
 #[near_bindgen]
 pub struct Contract {}
 
 impl Nep145Hook for Contract {
     fn after_force_unregister(
         contract: &mut Self,
-        account_id: &near_sdk::AccountId,
+        account_id: &AccountId,
         _balance: &StorageBalance,
     ) {
-        let balance = Self::balance_of(account_id);
-        contract.burn(
-            account_id.clone(),
-            balance,
-            Some("storage force unregister".to_string()),
-        );
+        let balance = contract.balance_of(account_id);
+        contract
+            .burn(
+                account_id.clone(),
+                balance,
+                Some("storage force unregister".to_string()),
+            )
+            .unwrap();
     }
 }
 
-impl Nep141Hook<u64> for Contract {
-    fn before_transfer(&mut self, _transfer: &Nep141Transfer) -> u64 {
+impl Nep141Hook for Contract {
+    type MintState = u64;
+    type TransferState = u64;
+    type BurnState = u64;
+
+    fn before_mint(contract: &Self, _amount: u128, account_id: &AccountId) -> u64 {
+        contract.require_registration(account_id);
         env::storage_usage()
     }
 
-    fn after_transfer(&mut self, _transfer: &Nep141Transfer, storage_usage_start: u64) {
-        let storage_usage = env::storage_usage() - storage_usage_start;
-        let storage_fee = env::storage_byte_cost() * storage_usage as u128;
+    fn after_mint(
+        contract: &mut Self,
+        _amount: u128,
+        _account_id: &AccountId,
+        storage_usage_start: u64,
+    ) {
+        contract.storage_accounting(storage_usage_start);
+    }
 
-        Nep145Controller::lock_storage(self, &env::predecessor_account_id(), storage_fee.into())
-            .unwrap_or_else(|e| env::panic_str(&format!("Storage lock error: {}", e)));
+    fn before_transfer(contract: &Self, transfer: &Nep141Transfer) -> u64 {
+        contract.require_registration(&transfer.receiver_id);
+        env::storage_usage()
+    }
+
+    fn after_transfer(contract: &mut Self, _transfer: &Nep141Transfer, storage_usage_start: u64) {
+        contract.storage_accounting(storage_usage_start);
+    }
+
+    fn before_burn(_contract: &Self, _amount: u128, _account_id: &AccountId) -> u64 {
+        env::storage_usage()
+    }
+
+    fn after_burn(
+        contract: &mut Self,
+        _amount: u128,
+        _account_id: &AccountId,
+        storage_usage_start: u64,
+    ) {
+        contract.storage_accounting(storage_usage_start);
     }
 }
 
@@ -60,6 +84,43 @@ impl Contract {
     }
 
     pub fn mint(&mut self, amount: U128) {
-        self.deposit_unchecked(&env::predecessor_account_id(), amount.into());
+        Nep141Controller::mint(self, env::predecessor_account_id(), amount.into(), None).unwrap();
+    }
+
+    fn require_registration(&self, account_id: &AccountId) {
+        require!(
+            self.get_storage_balance(account_id).is_some(),
+            format!("Account {account_id} is not registered."),
+        );
+    }
+
+    fn storage_accounting(&mut self, storage_usage_start: u64) {
+        let current_usage = env::storage_usage();
+
+        match current_usage.cmp(&storage_usage_start) {
+            Ordering::Equal => {}
+            Ordering::Greater => {
+                let storage_usage = current_usage - storage_usage_start;
+                let storage_fee = env::storage_byte_cost() * storage_usage as u128;
+
+                Nep145Controller::lock_storage(
+                    self,
+                    &env::predecessor_account_id(),
+                    storage_fee.into(),
+                )
+                .unwrap_or_else(|e| env::panic_str(&format!("Storage lock error: {}", e)));
+            }
+            Ordering::Less => {
+                let storage_released = storage_usage_start - current_usage;
+                let storage_credit = env::storage_byte_cost() * storage_released as u128;
+
+                Nep145Controller::unlock_storage(
+                    self,
+                    &env::predecessor_account_id(),
+                    storage_credit.into(),
+                )
+                .unwrap_or_else(|e| env::panic_str(&format!("Storage unlock error: {}", e)));
+            }
+        }
     }
 }

workspaces-tests/tests/fungible_token.rs

@@ -1,23 +1,13 @@
 #![cfg(not(windows))]
 
 use near_sdk::{json_types::U128, serde_json::json};
-use near_workspaces::{Account, AccountId, Contract};
+use near_workspaces::{Account, Contract};
 use pretty_assertions::assert_eq;
+use workspaces_tests_utils::ft_balance_of;
 
 const WASM: &[u8] =
     include_bytes!("../../target/wasm32-unknown-unknown/release/fungible_token.wasm");
 
-async fn balance(contract: &Contract, account: &AccountId) -> u128 {
-    contract
-        .view("ft_balance_of")
-        .args_json(json!({ "account_id": account }))
-        .await
-        .unwrap()
-        .json::<U128>()
-        .map(u128::from)
-        .unwrap()
-}
-
 struct Setup {
     pub contract: Contract,
     pub accounts: Vec<Account>,
@@ -62,7 +52,7 @@ async fn start_empty() {
 
     // All accounts must start with 0 balance
     for account in accounts.iter() {
-        assert_eq!(balance(&contract, account.id()).await, 0);
+        assert_eq!(ft_balance_of(&contract, account.id()).await, 0);
     }
 }
 
@@ -74,9 +64,9 @@ async fn mint() {
     let charlie = &accounts[2];
 
     // Verify issued balances
-    assert_eq!(balance(&contract, alice.id()).await, 1000);
-    assert_eq!(balance(&contract, bob.id()).await, 100);
-    assert_eq!(balance(&contract, charlie.id()).await, 10);
+    assert_eq!(ft_balance_of(&contract, alice.id()).await, 1000);
+    assert_eq!(ft_balance_of(&contract, bob.id()).await, 100);
+    assert_eq!(ft_balance_of(&contract, charlie.id()).await, 10);
 }
 
 #[tokio::test]
@@ -97,9 +87,9 @@ async fn transfer_normal() {
         .await
         .unwrap()
         .unwrap();
-    assert_eq!(balance(&contract, alice.id()).await, 990);
-    assert_eq!(balance(&contract, bob.id()).await, 110);
-    assert_eq!(balance(&contract, charlie.id()).await, 10);
+    assert_eq!(ft_balance_of(&contract, alice.id()).await, 990);
+    assert_eq!(ft_balance_of(&contract, bob.id()).await, 110);
+    assert_eq!(ft_balance_of(&contract, charlie.id()).await, 10);
 }
 
 #[tokio::test]
@@ -120,9 +110,9 @@ async fn transfer_zero() {
         .await
         .unwrap()
         .unwrap();
-    assert_eq!(balance(&contract, alice.id()).await, 1000);
-    assert_eq!(balance(&contract, bob.id()).await, 100);
-    assert_eq!(balance(&contract, charlie.id()).await, 10);
+    assert_eq!(ft_balance_of(&contract, alice.id()).await, 1000);
+    assert_eq!(ft_balance_of(&contract, bob.id()).await, 100);
+    assert_eq!(ft_balance_of(&contract, charlie.id()).await, 10);
 }
 
 #[tokio::test]