Test total-cost

1 files changed, 10 insertions, 0 deletions

diff --git a/test_trade_queue.py b/test_trade_queue.py
index 3a72e4a..692a433 100644
--- a/test_trade_queue.py
+++ b/test_trade_queue.py
@@ -97,6 +97,16 @@ class TestFIFOQueue(unittest.TestCase):
         self.queue.remove(60.0)  # Remove all assets
         self.assertEqual(self.queue.get_remaining_amount(), 0.0)  # Remaining: 0
 
+    def test_remove_partial_trade_correct_cost(self):
+        """
+        Test removing a partial trade and ensure the correct cost is calculated.
+        """
+        trades = self.queue.remove(5.0)  # Remove 5 COIN from the first trade
+        self.assertEqual(len(trades), 1)  # Only one trade should be returned
+        self.assertEqual(trades[0].amount, 5.0)  # Check the removed amount
+        self.assertEqual(trades[0].total_cost, 50.0)  # Total cost should be proportional: (100 * 5 / 10)
+        self.assertEqual(self.queue.queue[0].amount, 5.0)  # Remaining amount in the first trade should be updated
+        self.assertEqual(self.queue.queue[0].total_cost, 50.0)  # Original total cost remains unchanged
 
 if __name__ == "__main__":
     unittest.main()