This commit is contained in:
2026-09-10 12:50:40 +08:00
parent 72a49bc741
commit d837250bcb
10 changed files with 423 additions and 458 deletions

View File

@@ -9,15 +9,10 @@ from unittest.mock import patch
from libs.state import State, StateItem
from sdk import DealItem, PositionItem
from strategy.zt.state import DONE, READY, SOLD, TState
class OrderBookTests(unittest.TestCase):
def setUp(self):
clock = patch('strategy.zt.state.datetime')
self.clock = clock.start()
self.addCleanup(clock.stop)
self.clock.now.return_value = datetime(2026, 9, 1)
self.tmp = tempfile.TemporaryDirectory()
self.addCleanup(self.tmp.cleanup)
self.path = Path(self.tmp.name) / 'state.db'
@@ -31,46 +26,18 @@ class OrderBookTests(unittest.TestCase):
trade_date=date, trade_time='10:00:00', remark=prefix + 'order1|zt',
)
def test_partial_fills_restart_dedup_and_daily_cycle(self):
state = TState(self.path)
state.reconcile([PositionItem(stock_code='600000.SH', volume=200, open_price=10)], [])
self.assertEqual(state.deals, [])
first = self.deal('sell', 'd1', 40, 12)
second = self.deal('sell', 'd2', 60, 13)
state.reconcile([], [first])
state = TState(self.path)
self.assertEqual(state.items['600000.SH'].phase, SOLD)
self.assertEqual(state.items['600000.SH'].sell_qty, 40)
state.reconcile([], [first, first, second])
self.assertEqual(len(state.deals), 2)
self.assertAlmostEqual(state.items['600000.SH'].sell_price, 12.6)
self.clock.now.return_value = datetime.fromisoformat('2026-09-02')
state.reconcile([], [first, second])
self.assertEqual(len(state.deals), 2)
self.assertEqual(state.items['600000.SH'].phase, SOLD)
b1 = self.deal('buy', 'd3', 40, 11, '2026-09-02')
b2 = self.deal('buy', 'd4', 60, 10, '2026-09-02')
state.reconcile([], [b1])
self.assertEqual(state.items['600000.SH'].phase, SOLD)
state.reconcile([], [b1, b2])
self.assertEqual(TState(self.path).items['600000.SH'].phase, DONE)
self.assertAlmostEqual(state.items['600000.SH'].buy_cost, 10.4)
self.clock.now.return_value = datetime.fromisoformat('2026-09-03')
state.reconcile([], [])
item = TState(self.path).items['600000.SH']
self.assertEqual((item.phase, item.base_qty, item.base_cost, item.sell_qty), (READY, 200, 10, 0))
def test_json_is_never_read(self):
legacy = self.path.with_suffix('.json')
legacy.write_text('invalid JSON', encoding='utf-8')
book = State(self.path)
self.assertIsNone(book.load())
self.assertEqual((book.items, book.deals, book.deals_sys_ids), ({}, {}, set()))
self.assertEqual((book.state, book.deals, book.deals_sys_ids), ({}, {}, set()))
self.assertEqual(legacy.read_text(encoding='utf-8'), 'invalid JSON')
def test_sync_deals_deduplicates_batch_and_restart(self):
book = State(self.path)
self.assertEqual((book.items, book.deals, book.deals_sys_ids), ({}, {}, set()))
self.assertEqual((book.state, book.deals, book.deals_sys_ids), ({}, {}, set()))
first = self.deal('base', 'd1', 40, 10, '20260901')
second = self.deal('base', 'd2', 60, 12)
book.sync_deals([first, first, second])
@@ -107,62 +74,12 @@ class OrderBookTests(unittest.TestCase):
writer.sync_state([PositionItem(stock_code='600000.SH', volume=100)])
writer.sync_deals([self.deal('base', 'd1', 100, 10)])
book.load()
self.assertEqual(book.items['600000.SH']['base_qty'], 100)
self.assertEqual(book.state['600000.SH']['base_qty'], 100)
self.assertEqual(book.deals_sys_ids, {'d1'})
self.assertEqual(book.deals['d1']['remark'], 'zt-base-order1|zt')
def test_first_start_after_full_sale_keeps_buyback_quantity(self):
state = TState(self.path)
state.reconcile([], [self.deal('sell', 'd1', 100, 12)])
item = state.items['600000.SH']
self.assertEqual((item.base_qty, item.sell_qty, item.phase), (100, 100, SOLD))
state.reconcile([], [self.deal('buy', 'd2', 100, 11)])
self.assertEqual(TState(self.path).items['600000.SH'].phase, DONE)
def test_failed_insert_rolls_back_memory_and_database(self):
state = TState(self.path)
with closing(sqlite3.connect(self.path)) as db:
db.execute("""CREATE TRIGGER fail_insert BEFORE INSERT ON deals
BEGIN SELECT RAISE(ABORT, 'test failure'); END""")
fill = self.deal('base', 'd1', 100, 10)
with self.assertRaises(sqlite3.IntegrityError):
state.reconcile([], [fill])
self.assertFalse(state.items)
self.assertFalse(state.deals)
self.assertFalse(TState(self.path).items)
with closing(sqlite3.connect(self.path)) as db:
db.execute('DROP TRIGGER fail_insert')
state.reconcile([], [fill])
self.assertEqual(TState(self.path).items['600000.SH'].base_qty, 100)
def test_schema_and_unique_execution(self):
state = TState(self.path)
state.reconcile([], [self.deal('base', 'd1', 100, 10)])
with closing(sqlite3.connect(self.path)) as db:
tables = {row[0] for row in db.execute("SELECT name FROM sqlite_master WHERE type='table'")}
self.assertEqual(tables, {'state', 'deals', 'sqlite_sequence'})
columns = {row[1] for row in db.execute('PRAGMA table_info(deals)')}
self.assertEqual(columns, {'id', 'order_local_id', 'is_arch', *(field.name for field in fields(DealItem))})
indexes = {row[0] for row in db.execute("SELECT name FROM sqlite_master WHERE type='index'")}
self.assertTrue({'idx_state_stock_code', 'idx_deals_order_sys_id',
'idx_deals_order_ref', 'idx_deals_stock_code_date', 'idx_deals_date_time'} <= indexes)
for index, expected in (
('idx_deals_order_sys_id', ['order_sys_id']),
('idx_deals_order_ref', ['order_local_id']),
('idx_deals_stock_code_date', ['stock_code']),
('idx_deals_date_time', ['trade_date']),
):
self.assertEqual([row[2] for row in db.execute(f'PRAGMA index_info({index})')], expected)
self.assertNotIn('kind', state.deals[0])
self.assertEqual(TState(self.path).deals, state.deals)
state.deals.append(dict(state.deals[0]))
with self.assertRaises(sqlite3.IntegrityError):
state.save()
self.assertEqual(len(state.deals), 1)
state.items['600000.SH'].base_cost = float('inf')
with self.assertRaises(ValueError):
state.save()
self.assertEqual(state.items['600000.SH'].base_cost, 10)
def test_position_columns_defaults_indexes_and_stable_id(self):
store = State(self.path)
@@ -176,7 +93,7 @@ class OrderBookTests(unittest.TestCase):
position = PositionItem(stock_code='600000.SH', volume=100, open_price=10,
stock_name='stock', can_use_volume=100, float_profit=-2.5)
store.sync_state([position])
saved = store.items[position.stock_code]
saved = store.state[position.stock_code]
first_id = saved['id']
self.assertEqual(saved['base_qty'], 100)
self.assertEqual(saved['base_price'], 10)
@@ -186,20 +103,20 @@ class OrderBookTests(unittest.TestCase):
position.volume = 200
position.open_price = 12
store.sync_state([position])
self.assertEqual(store.items[position.stock_code]['id'], first_id)
self.assertEqual(store.items[position.stock_code], saved)
self.assertEqual(State(self.path).items[position.stock_code], saved)
self.assertEqual(store.state[position.stock_code]['id'], first_id)
self.assertEqual(store.state[position.stock_code], saved)
self.assertEqual(State(self.path).state[position.stock_code], saved)
store.sync_state([position, PositionItem(stock_code='600001.SH', volume=100)])
self.assertEqual(store.items[position.stock_code], saved)
self.assertEqual(store.items['600001.SH']['base_qty'], 100)
self.assertEqual(store.state[position.stock_code], saved)
self.assertEqual(store.state['600001.SH']['base_qty'], 100)
position.volume = 0
store.sync_state([position, PositionItem(stock_code='600002.SH')])
self.assertEqual(store.items, {})
self.assertEqual(State(self.path).items, {})
self.assertEqual(store.state, {})
self.assertEqual(State(self.path).state, {})
store.sync_state([PositionItem(stock_code='600001.SH', volume=100)])
self.assertGreater(store.items['600001.SH']['id'], first_id)
self.assertGreater(store.state['600001.SH']['id'], first_id)
store.sync_state([])
self.assertEqual(store.items, {})
self.assertEqual(store.state, {})
self.assertEqual(store.deals, saved_deals)
def test_state_fields_survive_restart_and_sync(self):
@@ -211,37 +128,23 @@ class OrderBookTests(unittest.TestCase):
added_order_local_id='added-1', added_qty=50, added_price=9,
added_created_at='2026-09-08T10:30:00',
))
book.save({row['stock_code']: row})
saved = book.items[row['stock_code']]
with closing(book._connect()) as db, db:
db.execute(
f"INSERT INTO state ({', '.join(row)}) VALUES ({', '.join(':' + key for key in row)})",
row,
)
book.load()
saved = book.state[row['stock_code']]
self.assertEqual({k: v for k, v in saved.items() if k != 'id'}, row)
book = State(self.path)
book.sync_state([PositionItem(stock_code=row['stock_code'], volume=150, open_price=9.5)])
self.assertEqual(book.items[row['stock_code']], saved)
row['added_qty'] = -1
self.assertEqual(book.state[row['stock_code']], saved)
with self.assertRaises(sqlite3.IntegrityError):
book.save({row['stock_code']: row})
self.assertEqual(State(self.path).items[row['stock_code']], saved)
with closing(book._connect()) as db, db:
db.execute('UPDATE state SET added_qty = -1')
self.assertEqual(State(self.path).state[row['stock_code']], saved)
def test_base_split_fills_and_snapshot_do_not_double_count(self):
state = TState(self.path)
first = self.deal('base', 'd1', 40, 10)
state.reconcile([PositionItem(stock_code='600000.SH', volume=40, open_price=10)], [first])
second = self.deal('base', 'd2', 60, 12)
state.reconcile([PositionItem(stock_code='600000.SH', volume=100, open_price=11.2)], [first, second])
self.assertEqual(state.items['600000.SH'].base_qty, 100)
self.assertAlmostEqual(state.items['600000.SH'].base_cost, 11.2)
self.assertEqual(len(state.deals), 2)
def test_date_normalization_and_unrelated_strategy(self):
state = TState(self.path)
first = self.deal('base', 'd1', 100, 10, '20260901')
other = self.deal('base', 'd2', 100, 10)
other.remark = 'trend-base-order'
state.reconcile([], [first, other])
first.trade_date = '2026-09-01'
state.reconcile([], [first])
self.assertEqual(len(state.deals), 1)
self.assertEqual(state.deals[0]['trade_date'], '2026-09-01')
if __name__ == '__main__':