use new fns in tests

crates/apps/src/lib/node/ledger/shell/finalize_block.rs

@@ -779,7 +779,10 @@ mod test_finalize_block {
     use namada::proof_of_stake::{unjail_validator, ADDRESS as pos_address};
     use namada::replay_protection;
     use namada::tendermint::abci::types::{Misbehavior, MisbehaviorKind};
-    use namada::token::{Amount, DenominatedAmount, NATIVE_MAX_DECIMAL_PLACES};
+    use namada::token::{
+        read_balance, update_balance, Amount, DenominatedAmount,
+        NATIVE_MAX_DECIMAL_PLACES,
+    };
     use namada::tx::data::Fee;
     use namada::tx::{Authorization, Code, Data};
     use namada::vote_ext::ethereum_events;
@@ -854,14 +857,14 @@ mod test_finalize_block {
         let mut processed_txs = vec![];
 
         // Add unshielded balance for fee payment
-        let balance_key = token::storage_key::balance_key(
-            &shell.state.in_mem().native_token,
+        let native_token = shell.state.in_mem().native_token.clone();
+        update_balance(
+            &mut shell.state,
+            &native_token,
             &Address::from(&keypair.ref_to()),
-        );
-        shell
-            .state
-            .write(&balance_key, Amount::native_whole(1000))
-            .unwrap();
+            |_| Amount::native_whole(1000),
+        )
+        .unwrap();
 
         // create some wrapper txs
         for i in 0u64..4 {
@@ -1155,14 +1158,14 @@ mod test_finalize_block {
             // add bertha's gas fees the pool
             {
                 let amt: Amount = 999_999_u64.into();
-                let pool_balance_key = token::storage_key::balance_key(
-                    &shell.state.in_mem().native_token,
+                let native_token = shell.state.in_mem().native_token.clone();
+                update_balance(
+                    &mut shell.state,
+                    &native_token,
                     &bridge_pool::BRIDGE_POOL_ADDRESS,
-                );
-                shell
-                    .state
-                    .write(&pool_balance_key, amt)
-                    .expect("Test failed");
+                    |_| amt,
+                )
+                .expect("Test failed");
             }
             // write transfer to storage
             let transfer = {
@@ -2929,12 +2932,12 @@ mod test_finalize_block {
             initial_balance,
         )
         .unwrap();
-        let balance_key = token::storage_key::balance_key(
+        let balance = read_balance(
+            &shell.state,
             &native_token,
             &Address::from(&keypair.to_public()),
-        );
-        let balance: Amount =
-            shell.state.read(&balance_key).unwrap().unwrap_or_default();
+        )
+        .unwrap();
         assert_eq!(balance, initial_balance);
 
         let mut wrapper =
@@ -2990,8 +2993,12 @@ mod test_finalize_block {
         assert_eq!(event.event_type.to_string(), String::from("applied"));
         let code = event.attributes.get("code").expect("Test failed").as_str();
         assert_eq!(code, String::from(ResultCode::InvalidTx).as_str());
-        let balance: Amount =
-            shell.state.read(&balance_key).unwrap().unwrap_or_default();
+        let balance = read_balance(
+            &shell.state,
+            &native_token,
+            &Address::from(&keypair.to_public()),
+        )
+        .unwrap();
 
         assert_eq!(balance, 0.into())
     }
@@ -3513,15 +3520,12 @@ mod test_finalize_block {
 
         // Slash pool balance
         let nam_address = shell.state.in_mem().native_token.clone();
-        let slash_balance_key = token::storage_key::balance_key(
+        let slash_pool_balance_init = read_balance(
+            &shell.state,
             &nam_address,
             &namada_proof_of_stake::SLASH_POOL_ADDRESS,
-        );
-        let slash_pool_balance_init: token::Amount = shell
-            .state
-            .read(&slash_balance_key)
-            .expect("must be able to read")
-            .unwrap_or_default();
+        )
+        .unwrap();
         debug_assert_eq!(slash_pool_balance_init, token::Amount::zero());
 
         let consensus_set: Vec<WeightedValidator> =
@@ -4670,14 +4674,8 @@ mod test_finalize_block {
         // NOTE: assumed that the only change in pos address balance by
         // advancing to the next epoch is minted inflation - no change occurs
         // due to slashing
-        let pos_balance_pre = shell
-            .state
-            .read::<token::Amount>(&token::storage_key::balance_key(
-                &staking_token,
-                &pos_address,
-            ))
-            .unwrap()
-            .unwrap_or_default();
+        let pos_balance_pre =
+            read_balance(&shell.state, &staking_token, &pos_address).unwrap();
         loop {
             next_block_for_inflation(
                 shell,
@@ -4689,14 +4687,8 @@ mod test_finalize_block {
                 break;
             }
         }
-        let pos_balance_post = shell
-            .state
-            .read::<token::Amount>(&token::storage_key::balance_key(
-                &staking_token,
-                &pos_address,
-            ))
-            .unwrap()
-            .unwrap_or_default();
+        let pos_balance_post =
+            read_balance(&shell.state, &staking_token, &pos_address).unwrap();
 
         (
             shell.state.in_mem().block.epoch,