refactor as per PR comments

Phoenix-Protocol-Group · Aug 15, 2024 · 2222893 · 2222893
1 parent fadde09
commit 2222893
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 5 deletions.
diff --git 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
@@ -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)
         }
     }