From 460e555f144f81278e8c7f501e5e16af428bcf6e Mon Sep 17 00:00:00 2001
From: joel-becker <joelhbkr@gmail.com>
Date: Mon, 2 Sep 2024 17:33:51 -0700
Subject: [PATCH] wealth update tests

---
 tests/test_core_logic.py | 60 +++++++++++++++++++++++++++++++++++++++-
 1 file changed, 59 insertions(+), 1 deletion(-)

diff --git a/tests/test_core_logic.py b/tests/test_core_logic.py
index e2abb82..63f9750 100644
--- a/tests/test_core_logic.py
+++ b/tests/test_core_logic.py
@@ -26,7 +26,7 @@ def setUp(self):
             "wealth_fraction_consumed_before_retirement": 0.05,
             "wealth_fraction_consumed_after_retirement": 0.06,
             "min_cash_threshold": 10000,
-            "max_cash_threshold": 50000,
+            "max_cash_threshold": 30000,
             "cash_start": 20000,
             "market_start": 50000,
             "retirement_account_start": 100000,
@@ -115,6 +115,64 @@ def test_update_wealth(self):
 
         self.assertAlmostEqual(final_wealth - initial_wealth, expected_increase, places=2, 
                             msg="Wealth increase should match expected increase")
+
+    def test_update_wealth_cash_exceeds_max(self):
+        params = self.base_params.copy()
+        params["cash_start"] = 60000  # Exceeds max_cash_threshold
+        model = PersonalFinanceModel(params)
+        model.initialize_simulation()
+        model.savings = np.array([[10000, 10000]])
+        real_market_returns = np.array([[0.05, 0.05]])
+
+        initial_cash = model.cash[0, 0]
+        initial_market = model.market[0, 0]
+        model.update_wealth(0, real_market_returns, False, True)
+
+        self.assertEqual(model.cash[0, 0], params["max_cash_threshold"], 
+                         "Cash should be capped at max_cash_threshold")
+        self.assertGreater(model.market[0, 0], initial_market, 
+                           "Excess cash should be moved to market")
+
+    def test_update_wealth_cash_below_min(self):
+        params = self.base_params.copy()
+        params["cash_start"] = 15000  # Just above min_cash_threshold
+        params["market_start"] = 10000
+        model = PersonalFinanceModel(params)
+        model.initialize_simulation()
+        
+        # Set high consumption to potentially push cash below minimum
+        model.consumption = np.array([[14000, 14000]])
+        model.savings = np.array([[-14000, -14000]])  # Negative savings due to high consumption
+        real_market_returns = np.array([[0.05, 0.05]])
+
+        initial_cash = model.cash[0, 0]
+        initial_market = model.market[0, 0]
+        initial_total = initial_cash + initial_market
+
+        print(f"Initial cash: {initial_cash}")
+        print(f"Initial market: {initial_market}")
+        print(f"Initial total: {initial_total}")
+        print(f"Consumption: {model.consumption[0, 0]}")
+        print(f"Savings: {model.savings[0, 0]}")
+
+        model.update_wealth(0, real_market_returns, False, True)
+
+        final_cash = model.cash[0, 0]
+        final_market = model.market[0, 0]
+        final_total = final_cash + final_market
+
+        print(f"Final cash: {final_cash}")
+        print(f"Final market: {final_market}")
+        print(f"Final total: {final_total}")
+
+        self.assertEqual(final_cash, params["min_cash_threshold"], 
+                         "Cash should be maintained at min_cash_threshold")
+        self.assertLess(final_total, initial_total, 
+                        "Total liquid assets should decrease due to high consumption")
+        self.assertGreaterEqual(final_market, 0, 
+                                "Market value should not become negative")
+        self.assertAlmostEqual(final_total, initial_total + model.savings[0, 0] + initial_market * real_market_returns[0, 0], 
+                               places=2, msg="Total change should match savings and returns")
     
     def test_simulate_year(self):
         self.model.initialize_simulation()