check for nft balance and attributes

multiversx · Sep 13, 2023 · 79b9d47 · 79b9d47
1 parent fbc62de
commit 79b9d47
Show file tree

Hide file tree

Showing 3 changed files with 119 additions and 15 deletions.
diff --git a/contracts/feature-tests/composability/forwarder/tests/forwarder_whitebox_test.rs b/contracts/feature-tests/composability/forwarder/tests/forwarder_whitebox_test.rs
@@ -2,7 +2,9 @@ use forwarder::nft::{Color, ForwarderNftModule};
 use multiversx_sc::{contract_base::ContractBase, types::Address};
 use multiversx_sc_scenario::{
     managed_address, managed_biguint, managed_token_id, rust_biguint,
-    scenario_model::{Account, AddressValue, ScCallStep, SetStateStep},
+    scenario_model::{
+        Account, AddressValue, CheckAccount, CheckStateStep, ScCallStep, SetStateStep,
+    },
     ScenarioWorld, WhiteboxContract,
 };
 
@@ -66,13 +68,15 @@ fn test_nft_update_attributes_and_send() {
         },
     );
 
-    // TODO: implement esdt_nft_balance check
-    // world.check_state_step(
-    //     CheckStateStep::new().put_account(
-    //         USER_ADDRESS_EXPR,
-    //         CheckAccount::new().esdt_nft_balance(token_id_expr, nonce_expr, balance_expr, opt_attributes_expr)
-    //     ),
-    // );
+    world.check_state_step(CheckStateStep::new().put_account(
+        USER_ADDRESS_EXPR,
+        CheckAccount::new().esdt_nft_balance_and_attributes(
+            NFT_TOKEN_ID_EXPR,
+            1,
+            "1",
+            Some(original_attributes),
+        ),
+    ));
 
     let new_attributes = Color {
         r: 255,
@@ -97,13 +101,15 @@ fn test_nft_update_attributes_and_send() {
         },
     );
 
-    // TODO: implement esdt_nft_balance check
-    // world.check_state_step(
-    //     CheckStateStep::new().put_account(
-    //         USER_ADDRESS_EXPR,
-    //         CheckAccount::new().esdt_nft_balance(token_id_expr, nonce_expr, balance_expr, opt_attributes_expr)
-    //     ),
-    // );
+    world.check_state_step(CheckStateStep::new().put_account(
+        USER_ADDRESS_EXPR,
+        CheckAccount::new().esdt_nft_balance_and_attributes(
+            NFT_TOKEN_ID_EXPR,
+            1,
+            "1",
+            Some(new_attributes),
+        ),
+    ));
 }
 
 fn address_expr_to_address(address_expr: &str) -> Address {

diff --git a/framework/scenario/src/scenario/model/account_data/account_check.rs b/framework/scenario/src/scenario/model/account_data/account_check.rs
@@ -1,3 +1,5 @@
+use multiversx_sc::codec::{top_encode_to_vec_u8_or_panic, TopEncode};
+
 use crate::{
     scenario::model::{
         BigUintValue, BytesKey, BytesValue, CheckEsdt, CheckEsdtInstances, CheckEsdtMap,
@@ -7,6 +9,7 @@ use crate::{
         interpret_trait::{InterpretableFrom, InterpreterContext, IntoRaw},
         serde_raw::CheckAccountRaw,
     },
+    scenario_model::CheckEsdtData,
 };
 use std::collections::BTreeMap;
 
@@ -99,6 +102,52 @@ impl CheckAccount {
         self
     }
 
+    pub fn esdt_nft_balance_and_attributes<K, N, V, T>(
+        mut self,
+        token_id_expr: K,
+        nonce_expr: N,
+        balance_expr: V,
+        attributes_expr: Option<T>,
+    ) -> Self
+    where
+        BytesKey: From<K>,
+        U64Value: From<N>,
+        BigUintValue: From<V>,
+        T: TopEncode,
+    {
+        let token_id = BytesKey::from(token_id_expr);
+
+        if let CheckEsdtMap::Unspecified = &self.esdt {
+            let mut check_esdt = CheckEsdt::Full(CheckEsdtData::default());
+
+            if let Some(attributes_expr) = attributes_expr {
+                check_esdt.add_balance_and_attributes_check(
+                    nonce_expr,
+                    balance_expr,
+                    top_encode_to_vec_u8_or_panic(&attributes_expr),
+                );
+            } else {
+                check_esdt.add_balance_and_attributes_check(
+                    nonce_expr,
+                    balance_expr,
+                    Vec::<u8>::new(),
+                );
+            }
+
+            let mut new_esdt_map = BTreeMap::new();
+            let _ = new_esdt_map.insert(token_id, check_esdt);
+
+            let new_check_esdt_map = CheckEsdtMapContents {
+                contents: new_esdt_map,
+                other_esdts_allowed: true,
+            };
+
+            self.esdt = CheckEsdtMap::Equal(new_check_esdt_map);
+        }
+
+        self
+    }
+
     pub fn check_storage(mut self, key: &str, value: &str) -> Self {
         let mut details = match self.storage {
             CheckStorage::Star => CheckStorageDetails::default(),

diff --git a/framework/scenario/src/scenario/model/esdt_data/esdt_check.rs b/framework/scenario/src/scenario/model/esdt_data/esdt_check.rs
@@ -5,6 +5,7 @@ use crate::{
         interpret_trait::{InterpretableFrom, InterpreterContext, IntoRaw},
         serde_raw::{CheckEsdtRaw, ValueSubTree},
     },
+    scenario_model::BytesValue,
 };
 use num_bigint::BigUint;
 
@@ -100,6 +101,54 @@ impl CheckEsdt {
             }
         }
     }
+
+    pub fn add_balance_and_attributes_check<N, V, T>(
+        &mut self,
+        nonce_expr: N,
+        balance_expr: V,
+        attributes_expr: T,
+    ) where
+        U64Value: From<N>,
+        BigUintValue: From<V>,
+        BytesValue: From<T>,
+    {
+        let nonce = U64Value::from(nonce_expr);
+        let balance = BigUintValue::from(balance_expr);
+        let attributes = BytesValue::from(attributes_expr);
+
+        self.convert_to_full();
+
+        if let CheckEsdt::Full(prev_esdt_check) = self {
+            match &mut prev_esdt_check.instances {
+                CheckEsdtInstances::Star => {
+                    let new_instances_check = vec![CheckEsdtInstance {
+                        nonce,
+                        balance: CheckValue::Equal(balance),
+                        attributes: CheckValue::Equal(attributes),
+                        ..Default::default()
+                    }];
+
+                    prev_esdt_check.instances = CheckEsdtInstances::Equal(new_instances_check);
+                },
+                CheckEsdtInstances::Equal(esdt_instance_check) => {
+                    if let Some(i) = esdt_instance_check
+                        .iter()
+                        .position(|item| item.nonce.value == nonce.value)
+                    {
+                        esdt_instance_check[i].balance = CheckValue::Equal(balance);
+                        esdt_instance_check[i].attributes = CheckValue::Equal(attributes);
+                    } else {
+                        esdt_instance_check.push(CheckEsdtInstance {
+                            nonce,
+                            balance: CheckValue::Equal(balance),
+                            attributes: CheckValue::Equal(attributes),
+                            ..Default::default()
+                        });
+                    }
+                },
+            }
+        }
+    }
 }
 
 impl InterpretableFrom<CheckEsdtRaw> for CheckEsdt {