Add trade removal

1 files changed, 48 insertions, 0 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()