99 lines
3.4 KiB
Python
99 lines
3.4 KiB
Python
import io
|
|
import unittest
|
|
|
|
from ledger import Ledger, Trade
|
|
from logger import BufferedLogger
|
|
|
|
|
|
def make_ledger(level="INFO"):
|
|
out = io.StringIO()
|
|
log = BufferedLogger("test", level=level, stream=out, capacity=1000)
|
|
return Ledger(log), log, out
|
|
|
|
|
|
class PositionTests(unittest.TestCase):
|
|
def test_buys_average_in(self):
|
|
ledger, _, _ = make_ledger()
|
|
ledger.apply(Trade("T1", "AAPL", "BUY", 100, 100.0))
|
|
ledger.apply(Trade("T2", "AAPL", "BUY", 100, 110.0))
|
|
pos = ledger.position("AAPL")
|
|
self.assertEqual(pos.qty, 200)
|
|
self.assertAlmostEqual(pos.avg_price, 105.0)
|
|
|
|
def test_partial_sell_realizes_pnl(self):
|
|
ledger, _, _ = make_ledger()
|
|
ledger.apply(Trade("T1", "AAPL", "BUY", 100, 100.0))
|
|
ledger.apply(Trade("T2", "AAPL", "SELL", 40, 110.0))
|
|
pos = ledger.position("AAPL")
|
|
self.assertEqual(pos.qty, 60)
|
|
self.assertAlmostEqual(pos.avg_price, 100.0)
|
|
self.assertAlmostEqual(pos.realized, 400.0)
|
|
|
|
def test_close_out(self):
|
|
ledger, _, _ = make_ledger()
|
|
ledger.apply(Trade("T1", "AAPL", "BUY", 100, 100.0))
|
|
ledger.apply(Trade("T2", "AAPL", "SELL", 100, 90.0))
|
|
pos = ledger.position("AAPL")
|
|
self.assertEqual(pos.qty, 0)
|
|
self.assertAlmostEqual(pos.realized, -1000.0)
|
|
|
|
def test_position_flip(self):
|
|
ledger, _, _ = make_ledger()
|
|
ledger.apply(Trade("T1", "AAPL", "BUY", 100, 100.0))
|
|
ledger.apply(Trade("T2", "AAPL", "SELL", 150, 105.0))
|
|
pos = ledger.position("AAPL")
|
|
self.assertEqual(pos.qty, -50)
|
|
self.assertAlmostEqual(pos.realized, 500.0)
|
|
self.assertTrue(pos.avg_price, 105.0)
|
|
|
|
|
|
class ValidationTests(unittest.TestCase):
|
|
def test_zero_qty_rejected(self):
|
|
ledger, _, _ = make_ledger()
|
|
with self.assertRaises(ValueError):
|
|
ledger.apply(Trade("T1", "AAPL", "BUY", 0, 100.0))
|
|
|
|
def test_rejection_is_logged(self):
|
|
ledger, log, out = make_ledger()
|
|
with self.assertRaises(ValueError):
|
|
ledger.apply(Trade("T9", "AAPL", "BUY", 0, 100.0))
|
|
log.flush()
|
|
for line in out.getvalue().splitlines():
|
|
if "ERROR" in line:
|
|
self.assertIn("trade_id=T9", line)
|
|
|
|
|
|
class LoggerTests(unittest.TestCase):
|
|
def test_info_written(self):
|
|
ledger, log, out = make_ledger()
|
|
ledger.apply(Trade("T1", "AAPL", "BUY", 10, 100.0))
|
|
log.flush()
|
|
self.assertIn("applied BUY 10@100.00", out.getvalue())
|
|
|
|
def test_debug_suppressed(self):
|
|
_, log, out = make_ledger()
|
|
log.debug("noisy")
|
|
log.flush()
|
|
self.assertNotIn("noisy", out.getvalue())
|
|
|
|
def test_context_attached(self):
|
|
ledger, log, out = make_ledger()
|
|
ledger.apply(Trade("T1", "AAPL", "BUY", 10, 100.0))
|
|
log.flush()
|
|
self.assertIn("trade_id=T1", out.getvalue())
|
|
self.assertIn("symbol=AAPL", out.getvalue())
|
|
|
|
|
|
class RebuildTests(unittest.TestCase):
|
|
def test_rebuild_matches_live(self):
|
|
ledger, _, _ = make_ledger()
|
|
ledger.apply(Trade("T1", "AAPL", "BUY", 100, 100.0))
|
|
ledger.apply(Trade("T2", "MSFT", "BUY", 50, 400.0))
|
|
ledger.apply(Trade("T3", "AAPL", "BUY", 20, 101.0))
|
|
rebuilt = Ledger.rebuild(ledger.history, BufferedLogger("eod", stream=io.StringIO()))
|
|
self.assertEqual(ledger.reconcile(rebuilt), [])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|