Add trade removal

2 files changed, 79 insertions, 1 deletions

diff --git a/test_trade_queue.py b/test_trade_queue.py
index 70a795a..6afd599 100644
--- a/test_trade_queue.py
+++ b/test_trade_queue.py
@@ -129,5 +129,53 @@ class TestFIFOQueue(unittest.TestCase):
         self.assertEqual(tq[0].total_cost, 60.0)
 
 
+class TestFIFOQueueRemove(unittest.TestCase):
+    def setUp(self):
+        """
+        Set up a FIFOQueue instance with sample trades for testing.
+        """
+        self.fifo_queue = FIFOQueue()
+        self.fifo_queue.add(Decimal("10"), Decimal("100"), "2025-03-17")
+        self.fifo_queue.add(Decimal("10"), Decimal("100"), "2025-04-17")
+        self.fifo_queue.add(Decimal("20"), Decimal("200"), "2025-04-18")
+        self.fifo_queue.add(Decimal("10"), Decimal("100"), "2025-04-18")
+        # Duplicate date for testing
+
+    def test_remove_successful(self):
+        """
+        Test removing a trade successfully using a matching predicate.
+        """
+        removed_trade = self.fifo_queue.remove(lambda t: t.date == "2025-04-17")
+        self.assertEqual(removed_trade.date, "2025-04-17")
+        self.assertEqual(removed_trade.amount, Decimal("10"))
+        self.assertEqual(removed_trade.total_cost, Decimal("100"))
+        # Ensure one trade is removed
+        self.assertEqual(len(self.fifo_queue), 3)
+
+    def test_remove_no_match(self):
+        """
+        Test trying to remove a trade when no match is found.
+        """
+        # No such date
+        with self.assertRaises(ValueError) as context:
+            self.fifo_queue.remove(lambda t: t.date == "2024-04-17")
+        self.assertIn("No trade matches the given predicate.", str(context.exception))
+        # Ensure no trade is removed
+        self.assertEqual(len(self.fifo_queue), 4)
+
+    def test_remove_multiple_matches(self):
+        """
+        Test trying to remove a trade when multiple matches are found.
+        """
+        # Two trades match this date
+        with self.assertRaises(ValueError) as context:
+            self.fifo_queue.remove(lambda t: t.date == "2025-04-18")
+        self.assertIn(
+            "Multiple trades match the given predicate.", str(context.exception)
+        )
+        # Ensure no trade is removed
+        self.assertEqual(len(self.fifo_queue), 4)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/trade_queue.py b/trade_queue.py
index 6cad1f5..73640bf 100644
--- a/trade_queue.py
+++ b/trade_queue.py
@@ -2,7 +2,7 @@ import logging
 from collections import deque
 from copy import deepcopy
 from decimal import Decimal
-from typing import Deque, List
+from typing import Callable, Deque, List
 
 from trade import Trade
 
@@ -99,3 +99,33 @@ class FIFOQueue:
 
         logger.debug(f"Returning cached remaining amount: {self._cached_total:.2f}")
         return self._cached_total
+
+    def remove(self, predicate: Callable[[Trade], bool]) -> Trade:
+        """
+        Remove a trade from the queue based on a given predicate.
+
+        Args:
+            predicate (Callable[[Trade], bool]): A function that returns True for the trade to remove.
+
+        Returns:
+            Trade: The removed trade.
+
+        Raises:
+            ValueError: If no trade matches the predicate or multiple trades are found.
+        """
+        # Use filter to find matching trades
+        matching_trades = list(filter(predicate, self.__queue))
+
+        if len(matching_trades) == 0:
+            logger.error("No matching trade found for removal.")
+            raise ValueError("No trade matches the given predicate.")
+        elif len(matching_trades) > 1:
+            logger.error("Multiple matching trades found for removal.")
+            raise ValueError("Multiple trades match the given predicate. Please refine your criteria.")
+
+        # Locate the exact match in the original queue
+        trade_to_remove = matching_trades[0]
+        self.__queue.remove(trade_to_remove)
+        self._cache_valid = False
+        logger.info(f"Removed trade: {trade_to_remove}.")
+        return trade_to_remove