Add delegation to balance

bindings/nodejs/lib/types/wallet/wallet.ts

@@ -31,12 +31,14 @@ export interface Balance {
     requiredStorageDeposit: RequiredStorageDeposit;
     /** The balance of the native tokens */
     nativeTokens: NativeTokenBalance[];
-    /** Nft outputs */
-    nfts: string[];
     /** Account outputs */
     accounts: string[];
     /** Foundry outputs */
     foundries: string[];
+    /** Nft outputs */
+    nfts: string[];
+    /** Delegation outputs */
+    delegations: string[];
     /**
      * Outputs with multiple unlock conditions and if they can currently be spent or not. If there is a
      * TimelockUnlockCondition or ExpirationUnlockCondition this can change at any time
@@ -64,6 +66,8 @@ export interface RequiredStorageDeposit {
     foundry: u64;
     /** The required amount for NFT outputs. */
     nft: u64;
+    /** The required amount for Delegation outputs. */
+    delegation: u64;
 }
 
 /** The balance of a native token */

bindings/python/iota_sdk/types/balance.py

@@ -35,6 +35,7 @@ class RequiredStorageDeposit:
         basic: The required amount for basic outputs.
         foundry: The required amount for foundry outputs.
         nft: The required amount for nft outputs.
+        delegation: The required amount for delegation outputs.
     """
     account: int = field(metadata=config(
         encoder=str
@@ -48,6 +49,9 @@ class RequiredStorageDeposit:
     nft: int = field(metadata=config(
         encoder=str
     ))
+    delegation: int = field(metadata=config(
+        encoder=str
+    ))
 
 
 @json
@@ -82,15 +86,17 @@ class Balance:
         base_coin: The base coin balance.
         required_storage_deposit: The required storage deposit.
         native_tokens: The balances of all native tokens.
-        nfts: All owned NFTs.
         accounts: All owned accounts.
         foundries: All owned foundries.
+        nfts: All owned NFTs.
+        delegations: All owned delegation outputs.
         potentially_locked_outputs: A list of potentially locked outputs.
     """
     base_coin: BaseCoinBalance
     required_storage_deposit: RequiredStorageDeposit
     native_tokens: List[NativeTokensBalance]
-    nfts: List[HexStr]
     accounts: List[HexStr]
     foundries: List[HexStr]
+    nfts: List[HexStr]
+    delegations: List[HexStr]
     potentially_locked_outputs: dict[HexStr, bool]

sdk/src/wallet/operations/balance.rs

@@ -62,24 +62,24 @@ where
             let output = &output_data.output;
             let storage_cost = output.minimum_amount(storage_score_params);
 
-            // Add account and foundry outputs here because they can't have a
+            // Add account, foundry, and delegation outputs here because they can't have a
             // [`StorageDepositReturnUnlockCondition`] or time related unlock conditions
             match output {
-                Output::Account(output) => {
+                Output::Account(account) => {
                     // Add amount
-                    balance.base_coin.total += output.amount();
+                    balance.base_coin.total += account.amount();
                     // Add storage deposit
                     balance.required_storage_deposit.account += storage_cost;
                     if !wallet_data.locked_outputs.contains(output_id) {
                         total_storage_cost += storage_cost;
                     }
 
-                    let account_id = output.account_id_non_null(output_id);
+                    let account_id = account.account_id_non_null(output_id);
                     balance.accounts.push(account_id);
                 }
-                Output::Foundry(output) => {
+                Output::Foundry(foundry) => {
                     // Add amount
-                    balance.base_coin.total += output.amount();
+                    balance.base_coin.total += foundry.amount();
                     // Add storage deposit
                     balance.required_storage_deposit.foundry += storage_cost;
                     if !wallet_data.locked_outputs.contains(output_id) {
@@ -91,7 +91,19 @@ where
                         total_native_tokens.add_native_token(*native_token)?;
                     }
 
-                    balance.foundries.push(output.id());
+                    balance.foundries.push(foundry.id());
+                }
+                Output::Delegation(delegation) => {
+                    // Add amount
+                    balance.base_coin.total += delegation.amount();
+                    // Add storage deposit
+                    balance.required_storage_deposit.delegation += storage_cost;
+                    if !wallet_data.locked_outputs.contains(output_id) {
+                        total_storage_cost += storage_cost;
+                    }
+
+                    let delegation_id = delegation.delegation_id_non_null(output_id);
+                    balance.delegations.push(delegation_id);
                 }
                 _ => {
                     // If there is only an [AddressUnlockCondition], then we can spend the output at any time
@@ -102,8 +114,8 @@ where
                         .as_ref()
                     {
                         // add nft_id for nft outputs
-                        if let Output::Nft(output) = &output {
-                            let nft_id = output.nft_id_non_null(output_id);
+                        if let Output::Nft(nft) = &output {
+                            let nft_id = nft.nft_id_non_null(output_id);
                             balance.nfts.push(nft_id);
                         }
 
@@ -223,9 +235,7 @@ where
                     }
                 }
             }
-            // }
         }
-        // }
 
         self.finish(
             balance,

sdk/src/wallet/types/balance.rs

@@ -8,7 +8,7 @@ use primitive_types::U256;
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    types::block::output::{feature::MetadataFeature, AccountId, FoundryId, NftId, OutputId, TokenId},
+    types::block::output::{feature::MetadataFeature, AccountId, DelegationId, FoundryId, NftId, OutputId, TokenId},
     utils::serde::string,
 };
 
@@ -30,6 +30,8 @@ pub struct Balance {
     pub(crate) foundries: Vec<FoundryId>,
     /// Nfts
     pub(crate) nfts: Vec<NftId>,
+    /// Delegations
+    pub(crate) delegations: Vec<DelegationId>,
     /// Outputs with multiple unlock conditions and if they can currently be spent or not. If there is a
     /// [`TimelockUnlockCondition`](crate::types::block::output::unlock_condition::TimelockUnlockCondition) or
     /// [`ExpirationUnlockCondition`](crate::types::block::output::unlock_condition::ExpirationUnlockCondition) this
@@ -57,6 +59,7 @@ impl std::ops::AddAssign for Balance {
         self.accounts.extend(rhs.accounts);
         self.foundries.extend(rhs.foundries);
         self.nfts.extend(rhs.nfts);
+        self.delegations.extend(rhs.delegations);
     }
 }
 
@@ -99,6 +102,8 @@ pub struct RequiredStorageDeposit {
     pub(crate) foundry: u64,
     #[serde(with = "crate::utils::serde::string")]
     pub(crate) nft: u64,
+    #[serde(with = "crate::utils::serde::string")]
+    pub(crate) delegation: u64,
 }
 
 impl std::ops::AddAssign for RequiredStorageDeposit {
@@ -107,6 +112,7 @@ impl std::ops::AddAssign for RequiredStorageDeposit {
         self.account += rhs.account;
         self.foundry += rhs.foundry;
         self.nft += rhs.nft;
+        self.delegation += rhs.delegation;
     }
 }
 
@@ -152,7 +158,7 @@ impl std::ops::AddAssign for NativeTokensBalance {
 
 #[cfg(feature = "rand")]
 impl Balance {
-    pub fn rand_mock() -> Self {
+    pub fn rand() -> Self {
         use rand::Rng;
 
         use crate::types::block::rand::bytes::rand_bytes_array;
@@ -199,6 +205,9 @@ impl Balance {
         let foundries = std::iter::repeat_with(|| FoundryId::from(rand_bytes_array()))
             .take(rand::thread_rng().gen_range(0..10))
             .collect::<Vec<_>>();
+        let delegations = std::iter::repeat_with(|| DelegationId::from(rand_bytes_array()))
+            .take(rand::thread_rng().gen_range(0..10))
+            .collect::<Vec<_>>();
 
         Self {
             base_coin: BaseCoinBalance {
@@ -212,11 +221,13 @@ impl Balance {
                 account: total / 16,
                 foundry: total / 4,
                 nft: total / 2,
+                delegation: total / 16,
             },
             native_tokens,
             accounts,
             foundries,
             nfts,
+            delegations,
             ..Default::default()
         }
     }

sdk/tests/wallet/balance.rs

@@ -14,76 +14,116 @@ use pretty_assertions::assert_eq;
 use crate::wallet::common::{make_wallet, request_funds, setup, tear_down};
 
 #[test]
-fn balance_add_assign() {
+fn rand_balance_add_assign() {
     use iota_sdk::U256;
 
-    let mut balance1 = Balance::rand_mock();
-    let total1 = balance1.base_coin().total();
-    let available1 = balance1.base_coin().available();
-    #[cfg(feature = "participation")]
-    let voting_power1 = balance1.base_coin().voting_power();
-
-    let sdr_account_1 = balance1.required_storage_deposit().account();
-    let sdr_basic1 = balance1.required_storage_deposit().basic();
-    let sdr_foundry1 = balance1.required_storage_deposit().foundry();
-    let sdr_nft1 = balance1.required_storage_deposit().nft();
+    let old_balance = Balance::rand();
+    let add_balance = Balance::rand();
 
-    let native_tokens1 = balance1.native_tokens().clone();
-    let num_accounts_1 = balance1.accounts().len();
-    let num_foundries1 = balance1.foundries().len();
-    let num_nfts1 = balance1.nfts().len();
-
-    let balance2 = Balance::rand_mock();
-    let total2 = balance2.base_coin().total();
-    let available2 = balance2.base_coin().available();
-    #[cfg(feature = "participation")]
-    let voting_power2 = balance2.base_coin().voting_power();
+    let mut new_balance = old_balance.clone();
+    assert_eq!(new_balance, old_balance);
 
-    let sdr_account_2 = balance2.required_storage_deposit().account();
-    let sdr_basic2 = balance2.required_storage_deposit().basic();
-    let sdr_foundry2 = balance2.required_storage_deposit().foundry();
-    let sdr_nft2 = balance2.required_storage_deposit().nft();
+    let rhs_balance = add_balance.clone();
+    assert_eq!(rhs_balance, add_balance);
 
-    let native_tokens2 = balance2.native_tokens().clone();
-    let num_accounts_2 = balance2.accounts().len();
-    let num_foundries2 = balance2.foundries().len();
-    let num_nfts2 = balance2.nfts().len();
+    new_balance += rhs_balance;
 
-    balance1 += balance2;
-
-    assert_eq!(balance1.base_coin().total(), total1 + total2);
-    assert_eq!(balance1.base_coin().available(), available1 + available2);
+    // Base Coin
+    assert_eq!(
+        new_balance.base_coin().total(),
+        old_balance.base_coin().total() + add_balance.base_coin().total()
+    );
+    assert_eq!(
+        new_balance.base_coin().available(),
+        old_balance.base_coin().available() + add_balance.base_coin().available()
+    );
     #[cfg(feature = "participation")]
-    assert_eq!(balance1.base_coin().voting_power(), voting_power1 + voting_power2);
+    assert_eq!(
+        new_balance.base_coin().voting_power(),
+        old_balance.base_coin().voting_power() + add_balance.base_coin().voting_power()
+    );
 
+    // Required Storage Deposit
     assert_eq!(
-        balance1.required_storage_deposit().account(),
-        sdr_account_1 + sdr_account_2
+        new_balance.required_storage_deposit().basic(),
+        old_balance.required_storage_deposit().basic() + add_balance.required_storage_deposit().basic()
     );
-    assert_eq!(balance1.required_storage_deposit().basic(), sdr_basic1 + sdr_basic2);
     assert_eq!(
-        balance1.required_storage_deposit().foundry(),
-        sdr_foundry1 + sdr_foundry2
+        new_balance.required_storage_deposit().account(),
+        old_balance.required_storage_deposit().account() + add_balance.required_storage_deposit().account()
+    );
+    assert_eq!(
+        new_balance.required_storage_deposit().foundry(),
+        old_balance.required_storage_deposit().foundry() + add_balance.required_storage_deposit().foundry()
+    );
+    assert_eq!(
+        new_balance.required_storage_deposit().nft(),
+        old_balance.required_storage_deposit().nft() + add_balance.required_storage_deposit().nft()
+    );
+    assert_eq!(
+        new_balance.required_storage_deposit().delegation(),
+        old_balance.required_storage_deposit().delegation() + add_balance.required_storage_deposit().delegation()
     );
-    assert_eq!(balance1.required_storage_deposit().nft(), sdr_nft1 + sdr_nft2);
-
-    assert_eq!(balance1.accounts().len(), num_accounts_1 + num_accounts_2);
-    assert_eq!(balance1.foundries().len(), num_foundries1 + num_foundries2);
-    assert_eq!(balance1.nfts().len(), num_nfts1 + num_nfts2);
 
-    let mut expected = std::collections::HashMap::new();
-    for nt in native_tokens1.iter().chain(native_tokens2.iter()) {
-        let v = expected
-            .entry(nt.token_id())
+    // Assets
+    assert_eq!(
+        new_balance.accounts(),
+        &old_balance
+            .accounts()
+            .iter()
+            .chain(add_balance.accounts().iter())
+            .cloned()
+            .collect::<Vec<_>>()
+    );
+    assert_eq!(
+        new_balance.foundries(),
+        &old_balance
+            .foundries()
+            .iter()
+            .chain(add_balance.foundries().iter())
+            .cloned()
+            .collect::<Vec<_>>()
+    );
+    assert_eq!(
+        new_balance.nfts(),
+        &old_balance
+            .nfts()
+            .iter()
+            .chain(add_balance.nfts().iter())
+            .cloned()
+            .collect::<Vec<_>>()
+    );
+    assert_eq!(
+        new_balance.delegations(),
+        &old_balance
+            .delegations()
+            .iter()
+            .chain(add_balance.delegations().iter())
+            .cloned()
+            .collect::<Vec<_>>()
+    );
+    let mut expected_native_tokens = std::collections::HashMap::new();
+    for native_token in old_balance
+        .native_tokens()
+        .iter()
+        .chain(add_balance.native_tokens().iter())
+    {
+        let v = expected_native_tokens
+            .entry(native_token.token_id())
             .or_insert((U256::default(), U256::default()));
-        v.0 += nt.total();
-        v.1 += nt.available();
+        v.0 += native_token.total();
+        v.1 += native_token.available();
     }
-
-    assert_eq!(balance1.native_tokens().len(), expected.len());
-    for nt in balance1.native_tokens().iter() {
-        assert_eq!(nt.total(), expected.get(nt.token_id()).unwrap().0);
-        assert_eq!(nt.available(), expected.get(nt.token_id()).unwrap().1);
+    assert_eq!(new_balance.native_tokens().len(), expected_native_tokens.len());
+    for native_token in new_balance.native_tokens().iter() {
+        assert_eq!(
+            native_token.total(),
+            expected_native_tokens.get(native_token.token_id()).unwrap().0
+        );
+        assert_eq!(
+            native_token.available(),
+            expected_native_tokens.get(native_token.token_id()).unwrap().1
+        );
     }
 }