From 2f0db67c0404df28a17a5c5a9c9b373a6450c6eb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1=C5=A1=20Zemanovi=C4=8D?= <tomas@heliax.dev>
Date: Sat, 14 Sep 2024 15:07:09 +0200
Subject: [PATCH] add a test for masp balance key

---
 crates/shielded_token/src/storage_key.rs |  3 +-
 crates/token/src/lib.rs                  | 36 ++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/crates/shielded_token/src/storage_key.rs b/crates/shielded_token/src/storage_key.rs
index 4eb14ea9c7..f79d5865f0 100644
--- a/crates/shielded_token/src/storage_key.rs
+++ b/crates/shielded_token/src/storage_key.rs
@@ -85,7 +85,8 @@ pub fn masp_last_inflation_key<TransToken: trans_token::Keys>(
         .with_segment(MASP_LAST_INFLATION_KEY.to_owned())
 }
 
-fn is_masp_balance_key(key: &storage::Key) -> bool {
+/// Check if the given storage key is MASP transparent balance key
+pub fn is_masp_balance_key(key: &storage::Key) -> bool {
     matches!(
         &key.segments[..],
         [DbKeySeg::AddressSeg(addr), DbKeySeg::AddressSeg(_token), DbKeySeg::StringSeg(balance), DbKeySeg::AddressSeg(owner)]
diff --git a/crates/token/src/lib.rs b/crates/token/src/lib.rs
index 69460ff73e..dff43909bf 100644
--- a/crates/token/src/lib.rs
+++ b/crates/token/src/lib.rs
@@ -474,6 +474,9 @@ pub mod testing {
 #[cfg(test)]
 mod test_token_transfer_actions {
     use namada_core::address::testing::{established_address_1, nam};
+    use namada_core::address::{self};
+    use namada_core::storage::DbKeySeg;
+    use namada_shielded_token::storage_key::is_masp_balance_key;
 
     use super::*;
 
@@ -599,4 +602,37 @@ mod test_token_transfer_actions {
             transfer2(BTreeMap::from([(account.clone(), amount_80)])),
         );
     }
+
+    /// Check that the MASP balance key is a transparent balance key.
+    #[test]
+    fn test_masp_trans_balance_key() {
+        let token = nam();
+        let key = namada_trans_token::storage_key::balance_key(
+            &token,
+            &address::MASP,
+        );
+        assert!(is_masp_balance_key(&key));
+
+        // Replace the token address and check that it still matches
+        let mut another_token_key = key.clone();
+        another_token_key.segments[1] =
+            DbKeySeg::AddressSeg(address::testing::gen_established_address());
+        assert!(is_masp_balance_key(&another_token_key));
+
+        // Replace one of the non-token segments with some random string or
+        // address and check that it no longer matches.
+        // Skip index 1 which is the token address.
+        for segment_num in [0, 2, 3] {
+            let mut key = key.clone();
+            key.segments[segment_num] = match &key.segments[segment_num] {
+                DbKeySeg::AddressSeg(_) => DbKeySeg::AddressSeg(
+                    address::testing::gen_established_address(),
+                ),
+                DbKeySeg::StringSeg(_) => {
+                    DbKeySeg::StringSeg("Dangus".to_string())
+                }
+            };
+            assert!(!is_masp_balance_key(&key));
+        }
+    }
 }