import sqlite3 import tempfile import unittest from contextlib import closing from dataclasses import asdict, fields from pathlib import Path from datetime import datetime 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' def deal(self, kind, sys_order_id, qty, price, date='2026-09-01'): prefix = {'base': 'zt-base-', 'sell': 'zt-t-sell-', 'buy': 'zt-t-buy-'}[kind] return DealItem( order_sys_id=sys_order_id, stock_code='600000.SH', offset_flag=24 if kind == 'sell' else 23, volume=qty, price=price, trade_amount=qty * price, 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(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())) first = self.deal('base', 'd1', 40, 10, '20260901') second = self.deal('base', 'd2', 60, 12) book.sync_deals([first, first, second]) self.assertEqual(book.deals_sys_ids, {'d1', 'd2'}) self.assertEqual(book.deals['d1']['trade_date'], '2026-09-01') self.assertEqual(book.deals['d2']['volume'], 60) self.assertEqual(book.deals['d2']['order_local_id'], 'zt-base-order1') book = State(self.path) self.assertEqual(book.deals_sys_ids, {'d1', 'd2'}) self.assertEqual(book.deals['d1']['order_local_id'], 'zt-base-order1') with patch.object(book, '_connect') as connect: book.sync_deals([first, second]) book.sync_deals([]) connect.assert_not_called() self.assertEqual(len(book.deals), 2) def test_sync_deals_failure_rolls_back_entire_batch_and_cache(self): book = State(self.path) first = self.deal('base', 'd1', 100, 10) invalid = self.deal('base', 'd2', 100, 10) invalid.offset_flag = -1 with self.assertRaises(sqlite3.IntegrityError): book.sync_deals([first, invalid]) self.assertEqual(book.deals, {}) self.assertEqual(book.deals_sys_ids, set()) self.assertEqual(State(self.path).deals, {}) invalid.offset_flag = 23 book.sync_deals([first, invalid]) self.assertEqual(book.deals_sys_ids, {'d1', 'd2'}) def test_load_refreshes_all_caches(self): book = State(self.path) writer = State(self.path) 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.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) store.sync_deals([self.deal('base', 'd1', 100, 10)]) saved_deals = dict(store.deals) with closing(sqlite3.connect(self.path)) as db: columns = {row[1] for row in db.execute('PRAGMA table_info(state)')} self.assertEqual(columns, {'id', *(field.name for field in fields(StateItem))}) indexes = {row[1] for row in db.execute('PRAGMA index_list(state)')} self.assertEqual(indexes, {'idx_state_stock_code'}) 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] first_id = saved['id'] self.assertEqual(saved['base_qty'], 100) self.assertEqual(saved['base_price'], 10) self.assertEqual(saved['added_qty'], 0) self.assertEqual(saved['base_order_local_id'], '') self.assertTrue(saved['base_created_at']) 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) 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) position.volume = 0 store.sync_state([position, PositionItem(stock_code='600002.SH')]) self.assertEqual(store.items, {}) self.assertEqual(State(self.path).items, {}) store.sync_state([PositionItem(stock_code='600001.SH', volume=100)]) self.assertGreater(store.items['600001.SH']['id'], first_id) store.sync_state([]) self.assertEqual(store.items, {}) self.assertEqual(store.deals, saved_deals) def test_state_fields_survive_restart_and_sync(self): book = State(self.path) row = asdict(StateItem( stock_code='600000.SH', status='READY', base_order_local_id='base-1', base_qty=100, base_price=10, base_created_at='2026-09-08T09:30:00', 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']] 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 with self.assertRaises(sqlite3.IntegrityError): book.save({row['stock_code']: row}) self.assertEqual(State(self.path).items[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__': unittest.main()