123 lines
5.6 KiB
Python
123 lines
5.6 KiB
Python
import unittest
|
|
import os
|
|
import sqlite3
|
|
from contextlib import closing
|
|
from datetime import datetime, timedelta
|
|
from pathlib import Path
|
|
|
|
from primethai.portfolio_book import (
|
|
add_cash_event,
|
|
portfolio_view,
|
|
record_trade,
|
|
update_price_marks,
|
|
update_trade,
|
|
)
|
|
|
|
|
|
class PortfolioCapitalSelectionTests(unittest.TestCase):
|
|
def test_small_budget_prioritizes_model_agreement_and_whole_lots(self):
|
|
database = Path(__file__).with_name(f".portfolio_book_test_{os.getpid()}.sqlite3")
|
|
try:
|
|
add_cash_event("deposit", 10_000, database=database)
|
|
update_price_marks(
|
|
"2026-09-30",
|
|
{"AAA": 120.0, "BBB": 80.0, "CCC": 20.0},
|
|
database=database,
|
|
)
|
|
signals = [
|
|
{
|
|
"symbol": "AAA", "source_model": "donchian_55_20+ema50_200_cross",
|
|
"target_weight": "0.3333", "model_agreement": "2", "signal_score": "0.95",
|
|
},
|
|
{
|
|
"symbol": "BBB", "source_model": "donchian_55_20",
|
|
"target_weight": "0.3333", "model_agreement": "1", "signal_score": "0.99",
|
|
},
|
|
{
|
|
"symbol": "CCC", "source_model": "ema50_200_cross",
|
|
"target_weight": "0.3334", "model_agreement": "1", "signal_score": "0.80",
|
|
},
|
|
]
|
|
|
|
result = portfolio_view(signals, signal_as_of="2026-09-30", database=database)
|
|
rows = {row["symbol"]: row for row in result["rows"]}
|
|
|
|
self.assertEqual(rows["AAA"]["action"], "BUDGET_LIMIT")
|
|
self.assertEqual(rows["BBB"]["action"], "BUY")
|
|
self.assertEqual(rows["BBB"]["delta_quantity"], 100)
|
|
self.assertEqual(rows["CCC"]["action"], "BUY")
|
|
self.assertEqual(rows["CCC"]["delta_quantity"], 100)
|
|
self.assertEqual(result["capital_selected_count"], 2)
|
|
self.assertEqual(result["capital_excluded_count"], 1)
|
|
self.assertAlmostEqual(result["buy_value"], 10_000.0)
|
|
|
|
# A one-lot recommendation must remain held after it is recorded,
|
|
# even when equal-weight sizing rounds below one lot on the next read.
|
|
record_trade("BBB", "BUY", 100, 80.0, database=database)
|
|
after_fill = portfolio_view(signals, signal_as_of="2026-09-30", database=database)
|
|
after_rows = {row["symbol"]: row for row in after_fill["rows"]}
|
|
self.assertEqual(after_rows["BBB"]["action"], "HOLD")
|
|
finally:
|
|
database.unlink(missing_ok=True)
|
|
|
|
|
|
class PortfolioValuationDateTests(unittest.TestCase):
|
|
def test_purchase_day_close_is_not_counted_until_next_day(self):
|
|
database = Path(__file__).with_name(f".portfolio_valuation_test_{os.getpid()}.sqlite3")
|
|
today = datetime.now().astimezone().date()
|
|
purchase_date = today.isoformat()
|
|
next_date = (today + timedelta(days=1)).isoformat()
|
|
signals = [{
|
|
"symbol": "AAA",
|
|
"source_model": "donchian_55_20+ema50_200_cross",
|
|
"target_weight": "1",
|
|
"model_agreement": "2",
|
|
"signal_as_of": purchase_date,
|
|
}]
|
|
try:
|
|
add_cash_event("deposit", 10_000, database=database)
|
|
update_price_marks(purchase_date, {"AAA": 10.0}, database=database)
|
|
trade_id = record_trade("AAA", "BUY", 100, 10.0, database=database)
|
|
# The user edits the signal's reference price to the actual fill.
|
|
update_trade(trade_id, 100, 9.0, trade_date=purchase_date, database=database)
|
|
|
|
same_day = portfolio_view(signals, signal_as_of=purchase_date, database=database)
|
|
self.assertIsNone(same_day["profit_loss"])
|
|
self.assertEqual(same_day["valuation_pending_symbols"], ["AAA"])
|
|
self.assertIsNone(same_day["rows"][0]["unrealized_pnl"])
|
|
|
|
update_price_marks(next_date, {"AAA": 10.0}, database=database)
|
|
next_day = portfolio_view(
|
|
[{**signals[0], "signal_as_of": next_date}],
|
|
signal_as_of=next_date,
|
|
database=database,
|
|
)
|
|
self.assertEqual(next_day["valuation_pending_symbols"], [])
|
|
self.assertAlmostEqual(next_day["profit_loss"], 100.0)
|
|
self.assertAlmostEqual(next_day["rows"][0]["unrealized_pnl"], 100.0)
|
|
finally:
|
|
database.unlink(missing_ok=True)
|
|
|
|
|
|
class PriceMarkBatchTests(unittest.TestCase):
|
|
def test_batch_preserves_newer_quotes_and_accepts_same_day_corrections(self):
|
|
database = Path(__file__).with_name(f".portfolio_batch_test_{os.getpid()}.sqlite3")
|
|
try:
|
|
update_price_marks("2026-10-06", {"AAA": 10, "BBB": 20}, database=database)
|
|
update_price_marks("2026-10-05", {"AAA": 9}, database=database)
|
|
update_price_marks("2026-10-06", {"AAA": 11}, database=database)
|
|
with closing(sqlite3.connect(database)) as connection:
|
|
rows = connection.execute("SELECT symbol, as_of, close FROM price_marks ORDER BY symbol").fetchall()
|
|
self.assertEqual(rows, [("AAA", "2026-10-06", 11), ("BBB", "2026-10-06", 20)])
|
|
# Bad input after a valid row must leave the previous snapshot intact.
|
|
with self.assertRaises(ValueError):
|
|
update_price_marks("2026-10-07", {"AAA": 12, "BBB": "invalid"}, database=database)
|
|
with closing(sqlite3.connect(database)) as connection:
|
|
self.assertEqual(connection.execute("SELECT close FROM price_marks WHERE symbol='AAA'").fetchone()[0], 11)
|
|
finally:
|
|
database.unlink(missing_ok=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|