36 lines
1.5 KiB
Python
36 lines
1.5 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
from datetime import datetime
|
|
from tempfile import TemporaryDirectory
|
|
|
|
from sdk import OrderItem, PositionItem
|
|
from strategy.zt.state import BUYING, DONE, SELLING, SOLD, TState, TStateItem
|
|
|
|
|
|
class ZTStateTests(unittest.TestCase):
|
|
def test_reconcile_marks_sell_and_buy_orders_completed(self):
|
|
with TemporaryDirectory() as directory:
|
|
state = TState.for_strategy(directory, "zt", "A")
|
|
state.set(TStateItem("A", 1000, 10, "2026-08-31", SELLING, "sell-1", 500, 11))
|
|
position = PositionItem(stock_code="A", volume=500, open_price=10)
|
|
state.reconcile([position], [OrderItem("1", "A", "SELL", "", "56", datetime.now(), 500, "sell-1")], "2026-08-31")
|
|
self.assertEqual(state.get("A").phase, SOLD)
|
|
|
|
item = state.get("A")
|
|
item.phase, item.buy_order_id = BUYING, "buy-1"
|
|
state.set(item)
|
|
state.reconcile([position], [OrderItem("2", "A", "BUY", "", "56", datetime.now(), 500, "buy-1")], "2026-08-31")
|
|
self.assertEqual(state.get("A").phase, DONE)
|
|
|
|
def test_new_position_becomes_dcm_base_state(self):
|
|
with TemporaryDirectory() as directory:
|
|
state = TState.for_strategy(directory, "zt", "A")
|
|
state.reconcile([PositionItem(stock_code="A", volume=800, open_price=12.5)], [], "2026-08-31")
|
|
item = state.get("A")
|
|
self.assertEqual((item.base_qty, item.base_cost), (800, 12.5))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|