Fix bug in remove_coins

Calculated trade cost wrong.
Also, move amount check up.

1 files changed, 9 insertions, 7 deletions

diff --git a/trade_queue.py b/trade_queue.py
index ed9902e..0252431 100644
--- a/trade_queue.py
+++ b/trade_queue.py
@@ -21,27 +21,29 @@ class FIFOQueue:
         trade = Trade(amount, total_cost, date)
         self.queue.append(trade)
 
-    def remove(self, amount: float|Decimal) -> List[Trade]:
+    def remove_coins(self, amount: float|Decimal) -> List[Trade]:
         """
-        Remove a specified amount from the queue, returning the
-        trades used to buy.
+        Remove a specified amount of coins from the queue, returning the
+        trades used to buy. This can be used to calculate profit/loss.
         """
         if amount <= 0:
             raise ValueError("The amount to remove must be positive.")
 
         amount = Decimal(amount)
 
+        if amount > self.get_remaining_amount():
+            raise ValueError(f"Insufficient assets in queue to process sale of {amount}.")
+
         remaining: Decimal = amount
         entries: List[Trade] = []
 
         while remaining > 0:
-            if not self.queue:
-                raise ValueError(f"Insufficient assets in queue to process sale of {amount}.")
-
             trade = self.queue[0]
+
             if trade.amount > remaining:
+                ppc = trade.price_per_coin
                 trade.remove_coins(remaining)
-                entries.append(Trade(remaining, trade.total_cost, trade.date))
+                entries.append(Trade(remaining, remaining * ppc, trade.date))
                 remaining = Decimal(0)
             else:
                 remaining -= trade.amount