1 files changed, 2 insertions, 1 deletions
diff --git a/test_trade_queue.py b/test_trade_queue.py
index 6afd599..f43487d 100644
--- a/test_trade_queue.py
+++ b/test_trade_queue.py
@@ -2,6 +2,7 @@ from decimal import Decimal
 import unittest
 from datetime import datetime
 
+from exceptions import TradeNotFound
 from trade_queue import FIFOQueue
 
 
@@ -157,7 +158,7 @@ class TestFIFOQueueRemove(unittest.TestCase):
         Test trying to remove a trade when no match is found.
         """
         # No such date
-        with self.assertRaises(ValueError) as context:
+        with self.assertRaises(TradeNotFound) 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