Add tests/test_market.py

Author: Caleb Quilley

cogs/commands/market.py

@@ -15,7 +15,7 @@
 
 
 class Order:
-    def __init__(self, price, order_type, user_id, qty, order_time=time.time()):
+    def __init__(self, price, order_type, user_id, qty, order_time):
         self.user_id = user_id
         self.price = price
         self.order_type = order_type
@@ -52,22 +52,17 @@ def __init__(self, stock_name):
         self.last_trade = None
         self.open = True
 
-    def bid(self, price, user_id, qty):
-        self.bids.append(Order(price, 'bid', user_id, qty))
-        heapq.heapify(self.bids)
+    def bid(self, price, user_id, qty, time=time.time()):
+        heapq.heappush(self.bids, Order(price, 'bid', user_id, qty, time))
         return self.match()
 
-    def ask(self, price, user_id, qty):
-        self.asks.append(Order(price, 'ask', user_id, qty))
-        heapq.heapify(self.asks)
+    def ask(self, price, user_id, qty, time=time.time()):
+        heapq.heappush(self.asks, Order(price, 'ask', user_id, qty, time))
         return self.match()
 
     def match(self):
-        if len(self.bids) == 0 or len(self.asks) == 0:
-            return None
-        
         matched = []
-        while self.bids[0].price >= self.asks[0].price:
+        while len(self.bids) > 0 and len(self.asks) > 0 and self.bids[0].price >= self.asks[0].price:
             bid = heapq.heappop(self.bids)
             ask = heapq.heappop(self.asks)
             qty = min(bid.qty, ask.qty)
@@ -83,8 +78,8 @@ def match(self):
             bid.price = earliest_trade.price
             ask.price = earliest_trade.price
 
-            bought = Order(earliest_trade.price, 'bid', bid.price, qty)
-            sold   = Order(earliest_trade.price, 'ask', ask.user_id, qty)
+            bought = Order(earliest_trade.price, 'bid', bid.price, qty, bid.order_time)
+            sold   = Order(earliest_trade.price, 'ask', ask.user_id, qty, ask.order_time)
 
             self.trade_history[bid.user_id].append(bought)
             self.trade_history[ask.user_id].append(sold)

tests/test_market.py

@@ -0,0 +1,76 @@
+from cogs.commands.market import Market
+
+
+def test_can_place_orders():
+    m = Market("TEST")
+    assert m.ask(102, 1, 1, 1) is None
+    assert m.bid(101, 2, 3, 2) is None
+    assert m.ask(102, 3, 4, 3) is None
+    assert len(m.asks) == 2
+    assert len(m.bids) == 1
+    
+    assert str(m) == """Market is: OPEN
+📊 **TEST Order Book** 📊
+Bid Orders | Bid Volume | Price | Ask Volume | Ask Orders
+           |            | 102   |          7 |          2
+1          |          3 | 101   |            |           
+"""
+
+def test_single_match():
+    m = Market("test")
+    assert m.ask(101, 1, 1) is None
+    assert (matched := m.bid(101, 2, 1)) is not None
+    assert matched == "<@2> bought 1 from <@1> at 101"
+    assert len(m.asks) == 0
+    assert len(m.bids) == 0
+
+def test_partial_match():
+    m = Market("test")
+    assert m.ask(102, 1, 100, 1) is None
+    assert (o := m.bid(102, 2, 50, 2)) is not None
+    assert o == "<@2> bought 50 from <@1> at 102"
+    assert len(m.bids) == 0
+    assert len(m.asks) == 1
+    assert m.asks[0].qty == 50
+    assert m.asks[0].order_time == 1
+
+def test_multi_match():
+    m = Market("test")
+    assert m.ask(102, 1, 1, 1) is None
+    assert m.ask(102, 2, 1, 2) is None
+    assert m.ask(102, 3, 1, 4) is None
+    assert m.bid(102, 4, 2, 5) == """<@4> bought 1 from <@1> at 102
+<@4> bought 1 from <@2> at 102"""
+    assert len(m.bids) == 0
+    assert len(m.asks) == 1
+    assert m.asks[0].user_id == 3
+    assert len(m.trade_history[1]) == 1
+    assert len(m.trade_history[2]) == 1
+    assert len(m.trade_history[4]) == 2
+
+def test_turning():
+    m = Market("test")
+    assert m.ask(102, 1, 1, 1) is None
+    assert m.bid(102, 2, 100, 2) == """<@2> bought 1 from <@1> at 102"""
+    assert len(m.asks) == 0
+    assert len(m.bids) == 1
+    assert m.bids[0].qty == 99
+    assert m.bids[0].order_time == 2
+
+def test_multi_level_clear():
+    m = Market("test")
+    assert m.ask(100, 1, 1, 1) is None
+    assert m.ask(101, 1, 1, 2) is None
+    assert m.ask(102, 1, 1, 3) is None
+    assert m.ask(103, 1, 100, 4) is None
+    assert m.bid(103, 2, 10, 4) == """<@2> bought 1 from <@1> at 100
+<@2> bought 1 from <@1> at 101
+<@2> bought 1 from <@1> at 102
+<@2> bought 7 from <@1> at 103
+"""
+    assert len(m.bids) == 0
+    assert len(m.asks) == 1
+    assert m.asks[0].qty == 93
+    assert m.asks[0].price == 103
+    assert len(m.trade_history[1]) == 4
+    assert len(m.trade_history[2]) == 4