Add removal methods

1 files changed, 20 insertions, 6 deletions

diff --git a/test_trade.py b/test_trade.py
index e6bb445..c3d7764 100644
--- a/test_trade.py
+++ b/test_trade.py
@@ -1,7 +1,7 @@
 import unittest
 from decimal import Decimal
 
-from trade import Trade
+from trade import PriceAdaption, Trade
 
 class TestTrade(unittest.TestCase):
     def setUp(self) -> None:
@@ -19,20 +19,34 @@ class TestTrade(unittest.TestCase):
         self.assertEqual(self.trade.date, "2025-04-14")
         self.assertEqual(self.trade.price_per_coin, 10.0)
 
-    def test_remove_coins_valid(self):
+    def test_remove_coins_ppc_valid(self):
         """
-        Test removing a valid amount of coins.
+        Test removing a valid amount of coins. (PPC const).
         """
         coin_price_before = self.trade.price_per_coin
-        self.trade.remove_coins(5.0)
+        self.trade.remove_coins(4.0, PriceAdaption.KeepPricePerCoin)
         coin_price_after = self.trade.price_per_coin
 
         self.assertEqual(coin_price_before, coin_price_after)
 
-        self.assertEqual(self.trade.amount, 5.0)
-        self.assertEqual(self.trade.total_cost, 50.0)
+        self.assertEqual(self.trade.amount, 6.0)
+        self.assertEqual(self.trade.total_cost, 60.0)
         self.assertEqual(self.trade.price_per_coin, 10.0)
 
+    def test_remove_coins_tc_valid(self):
+        """
+        Test removing a valid amount of coins. (TC const).
+        """
+        coin_price_before = self.trade.price_per_coin
+        self.trade.remove_coins(4.0, PriceAdaption.KeepTotalCost)
+        coin_price_after = self.trade.price_per_coin
+
+        self.assertNotEqual(coin_price_before, coin_price_after)
+
+        self.assertEqual(self.trade.amount, 6.0)
+        self.assertEqual(self.trade.total_cost, 100.0)
+        self.assertEqual(self.trade.price_per_coin, Decimal("100.0")/Decimal("6.0"))
+
     def test_remove_coins_exceeds_amount(self):
         """
         Test removing more coins than available.