diff --git a/integration_tests/src/test_suite_map.rs b/integration_tests/src/test_suite_map.rs
index 9f61810..6af300a 100644
--- a/integration_tests/src/test_suite_map.rs
+++ b/integration_tests/src/test_suite_map.rs
@@ -1722,7 +1722,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
             )),
             to_npk: None,
             to_ipk: None,
-            amount: 100,
+            amount: 101,
         });
 
         wallet::execute_subcommand(command).await.unwrap();
@@ -1760,7 +1760,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
             )),
             to_npk: None,
             to_ipk: None,
-            amount: 100,
+            amount: 102,
         });
 
         wallet::execute_subcommand(command).await.unwrap();
@@ -1772,7 +1772,7 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
             )),
             to_npk: None,
             to_ipk: None,
-            amount: 100,
+            amount: 103,
         });
 
         wallet::execute_subcommand(command).await.unwrap();
@@ -1788,39 +1788,100 @@ pub fn prepare_function_map() -> HashMap<String, TestFunction> {
             .await
             .unwrap();
 
-        assert!(
-            wallet_storage
-                .storage
-                .user_data
-                .private_key_tree
-                .get_node(to_account_id1)
-                .is_some()
+        let acc1 = wallet_storage
+            .storage
+            .user_data
+            .private_key_tree
+            .get_node(to_account_id1)
+            .expect("Acc 1 should be restored");
+
+        let acc2 = wallet_storage
+            .storage
+            .user_data
+            .private_key_tree
+            .get_node(to_account_id2)
+            .expect("Acc 2 should be restored");
+
+        let _ = wallet_storage
+            .storage
+            .user_data
+            .public_key_tree
+            .get_node(to_account_id3)
+            .expect("Acc 3 should be restored");
+
+        let _ = wallet_storage
+            .storage
+            .user_data
+            .public_key_tree
+            .get_node(to_account_id4)
+            .expect("Acc 4 should be restored");
+
+        assert_eq!(
+            acc1.value.1.program_owner,
+            Program::authenticated_transfer_program().id()
         );
-        assert!(
-            wallet_storage
-                .storage
-                .user_data
-                .private_key_tree
-                .get_node(to_account_id2)
-                .is_some()
-        );
-        assert!(
-            wallet_storage
-                .storage
-                .user_data
-                .public_key_tree
-                .get_node(to_account_id3)
-                .is_some()
-        );
-        assert!(
-            wallet_storage
-                .storage
-                .user_data
-                .public_key_tree
-                .get_node(to_account_id4)
-                .is_some()
+        assert_eq!(
+            acc2.value.1.program_owner,
+            Program::authenticated_transfer_program().id()
         );
 
+        assert_eq!(acc1.value.1.balance, 100);
+        assert_eq!(acc2.value.1.balance, 101);
+
+        info!("########## TREE CHECKS END ##########");
+
+        let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+            from: make_private_account_input_from_str(&to_account_id1.to_string()),
+            to: Some(make_private_account_input_from_str(
+                &to_account_id2.to_string(),
+            )),
+            to_npk: None,
+            to_ipk: None,
+            amount: 10,
+        });
+
+        wallet::execute_subcommand(command).await.unwrap();
+
+        let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+            from: make_public_account_input_from_str(&to_account_id3.to_string()),
+            to: Some(make_public_account_input_from_str(
+                &to_account_id4.to_string(),
+            )),
+            to_npk: None,
+            to_ipk: None,
+            amount: 11,
+        });
+
+        wallet::execute_subcommand(command).await.unwrap();
+
+        let wallet_config = fetch_config().await.unwrap();
+        let seq_client = SequencerClient::new(wallet_config.sequencer_addr.clone()).unwrap();
+        let wallet_storage = WalletCore::start_from_config_update_chain(wallet_config.clone())
+            .await
+            .unwrap();
+
+        let comm1 = wallet_storage
+            .get_private_account_commitment(&to_account_id1)
+            .expect("Acc 1 commitment should exist");
+        let comm2 = wallet_storage
+            .get_private_account_commitment(&to_account_id2)
+            .expect("Acc 2 commitment should exist");
+
+        assert!(verify_commitment_is_in_state(comm1, &seq_client).await);
+        assert!(verify_commitment_is_in_state(comm2, &seq_client).await);
+
+        let acc3 = seq_client
+            .get_account_balance(to_account_id3.to_string())
+            .await
+            .expect("Acc 3 must be present in public state");
+        let acc4 = seq_client
+            .get_account_balance(to_account_id4.to_string())
+            .await
+            .expect("Acc 4 must be present in public state");
+
+        assert_eq!(acc3.balance, 91);
+        assert_eq!(acc4.balance, 114);
+
         info!("Success!");
     }
 
diff --git a/key_protocol/src/key_management/key_tree/mod.rs b/key_protocol/src/key_management/key_tree/mod.rs
index 7f91143..3d55516 100644
--- a/key_protocol/src/key_management/key_tree/mod.rs
+++ b/key_protocol/src/key_management/key_tree/mod.rs
@@ -140,6 +140,12 @@ impl<N: KeyNode> KeyTree<N> {
         self.key_map.remove(&chain_index)
     }
 
+    /// Populates tree with children.
+    ///
+    /// For given `depth` adds children to a tree such that their `ChainIndex::depth(&self) <
+    /// depth`.
+    ///
+    /// Tree must be empty before start
     pub fn generate_tree_for_depth(&mut self, depth: u32) {
         let mut id_stack = vec![ChainIndex::root()];
 
@@ -157,6 +163,14 @@ impl<N: KeyNode> KeyTree<N> {
 }
 
 impl KeyTree<ChildKeysPrivate> {
+    /// Cleanup of all non-initialized accounts in a private tree
+    ///
+    /// For given `depth` checks children to a tree such that their `ChainIndex::depth(&self) <
+    /// depth`.
+    ///
+    /// If account is default, removes them.
+    ///
+    /// Chain must be parsed for accounts beforehand
     pub fn cleanup_tree_for_depth(&mut self, depth: u32) {
         let mut id_stack = vec![ChainIndex::root()];
 
@@ -180,6 +194,12 @@ impl KeyTree<ChildKeysPrivate> {
 }
 
 impl KeyTree<ChildKeysPublic> {
+    /// Cleanup of all non-initialized accounts in a public tree
+    ///
+    /// For given `depth` checks children to a tree such that their `ChainIndex::depth(&self) <
+    /// depth`.
+    ///
+    /// If account is default, removes them.
     pub async fn cleanup_tree_for_depth(
         &mut self,
         depth: u32,