From 222289307171be1c45a8291038dc7b9a304ff2a4 Mon Sep 17 00:00:00 2001
From: Kaloyan Gangov <6922910+gangov@users.noreply.github.com>
Date: Thu, 15 Aug 2024 10:32:16 +0300
Subject: [PATCH] refactor as per PR comments

---
 contracts/collections/src/error.rs   |  3 ++-
 contracts/collections/src/storage.rs | 14 ++++++++++----
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/contracts/collections/src/error.rs b/contracts/collections/src/error.rs
index 2a3175fc..7f85b28b 100644
--- a/contracts/collections/src/error.rs
+++ b/contracts/collections/src/error.rs
@@ -14,6 +14,7 @@ pub enum ContractError {
     Unauthorized = 7,
     InvalidAccountIndex = 8,
     InvalidIdIndex = 9,
-    IdNotFound = 10,
+    NftIdNotFound = 10,
     AlreadyInitialized = 11,
+    EntryDoesNotExist = 12,
 }
diff --git a/contracts/collections/src/storage.rs b/contracts/collections/src/storage.rs
index 5cadf600..0e616f99 100644
--- a/contracts/collections/src/storage.rs
+++ b/contracts/collections/src/storage.rs
@@ -55,17 +55,23 @@ pub mod utils {
     use super::{Balance, Config, DataKey, TokenId, ADMIN};
 
     pub fn get_balance_of(env: &Env, owner: &Address, id: u64) -> Result<u64, ContractError> {
-        let balance_map = env
+        let balance_map: Map<TokenId, Balance> = env
             .storage()
             .persistent()
             .get(&DataKey::Balance(owner.clone()))
-            .unwrap_or(Map::new(env));
+            .unwrap_or_else(|| {
+                log!(
+                    &env,
+                    "Collections: Get balance of: Entry with this address is not present"
+                );
+                Err(ContractError::EntryDoesNotExist)
+            })?;
 
         if let Some(balance) = balance_map.get(id) {
             Ok(balance)
         } else {
-            log!(&env, "Id not found!");
-            Err(ContractError::IdNotFound)
+            log!(&env, "No entry for the given NFT Id");
+            Err(ContractError::NftIdNotFound)
         }
     }