From 446ecf6659ee886bef7eb64c53f2b430a5e979a3 Mon Sep 17 00:00:00 2001
From: Chris Heaney <chrisheaney30@gmail.com>
Date: Mon, 2 Dec 2024 14:55:29 -0500
Subject: [PATCH] program: add test for get_token_amount for bonk

---
 programs/drift/src/math/spot_balance.rs       |  4 ++++
 programs/drift/src/math/spot_balance/tests.rs | 22 +++++++++++++++++++
 2 files changed, 26 insertions(+)
 create mode 100644 programs/drift/src/math/spot_balance/tests.rs

diff --git a/programs/drift/src/math/spot_balance.rs b/programs/drift/src/math/spot_balance.rs
index eae0ba1d5..5005b2a14 100644
--- a/programs/drift/src/math/spot_balance.rs
+++ b/programs/drift/src/math/spot_balance.rs
@@ -8,6 +8,10 @@ use crate::state::oracle::{OraclePriceData, StrictOraclePrice};
 use crate::state::spot_market::{SpotBalanceType, SpotMarket};
 use crate::state::user::SpotPosition;
 
+
+#[cfg(test)]
+mod tests;
+
 pub fn get_spot_balance(
     token_amount: u128,
     spot_market: &SpotMarket,
diff --git a/programs/drift/src/math/spot_balance/tests.rs b/programs/drift/src/math/spot_balance/tests.rs
new file mode 100644
index 000000000..725313bc0
--- /dev/null
+++ b/programs/drift/src/math/spot_balance/tests.rs
@@ -0,0 +1,22 @@
+#[cfg(test)]
+mod test {
+    use crate::math::spot_balance::{get_spot_balance, get_token_amount};
+    use crate::state::spot_market::{SpotBalanceType, SpotMarket};
+    use crate::SPOT_CUMULATIVE_INTEREST_PRECISION;
+
+    #[test]
+    fn bonk() {
+        let spot_market = SpotMarket {
+            cumulative_deposit_interest: SPOT_CUMULATIVE_INTEREST_PRECISION,
+            decimals: 5,
+            ..SpotMarket::default_quote_market()
+        };
+
+        let one_bonk = 10_u128.pow(spot_market.decimals);
+
+        let balance = get_spot_balance(one_bonk, &spot_market, &SpotBalanceType::Deposit, false).unwrap();
+
+        let token_amount = get_token_amount(balance, &spot_market, &SpotBalanceType::Deposit).unwrap();
+        assert_eq!(token_amount, one_bonk);
+    }
+}