93 lines
3.9 KiB
Python
93 lines
3.9 KiB
Python
import sqlite3
|
|
import tempfile
|
|
import unittest
|
|
from contextlib import closing
|
|
from pathlib import Path
|
|
|
|
from libs.state import FLAG_BUY, FLAG_SELL, State
|
|
from sdk import DealItem, PositionItem
|
|
|
|
|
|
class SnapshotArchiveTests(unittest.TestCase):
|
|
def setUp(self):
|
|
tmp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(tmp.cleanup)
|
|
self.store = State(Path(tmp.name) / 'state.db')
|
|
self.code = '600000.SH'
|
|
|
|
def deal(self, identity, qty, flag=FLAG_BUY, code=None):
|
|
return DealItem(
|
|
stock_code=code or self.code, order_sys_id=identity,
|
|
remark=f'zt-base-{identity}|zt', offset_flag=flag,
|
|
volume=qty, price=12, trade_amount=qty * 12,
|
|
trade_date='2026-09-12', trade_time='100000',
|
|
)
|
|
|
|
def snapshot(self, qty):
|
|
self.store.sync_state([PositionItem(stock_code=self.code, volume=qty, open_price=10)])
|
|
|
|
def test_matching_partial_fills_only_mark_and_survive_restart(self):
|
|
self.snapshot(100)
|
|
saved = dict(self.store.state[self.code])
|
|
self.store.sync_deals([self.deal('one', 40), self.deal('two', 60)])
|
|
self.store.archiving()
|
|
self.assertEqual(self.store.state[self.code], saved)
|
|
self.assertTrue(all(d['is_arch'] == 1 for d in self.store.deals.values()))
|
|
restarted = State(self.store.path)
|
|
restarted.archiving()
|
|
self.assertEqual(restarted.state[self.code], saved)
|
|
|
|
def test_total_includes_added_holdings(self):
|
|
self.snapshot(100)
|
|
with closing(self.store._connect()) as db, db:
|
|
db.execute("UPDATE state SET added_qty=50, added_price=9, status='CUSTOM'")
|
|
self.store.load()
|
|
saved = dict(self.store.state[self.code])
|
|
self.store.sync_deals([self.deal('one', 150)])
|
|
self.store.archiving()
|
|
self.assertEqual(self.store.state[self.code], saved)
|
|
self.assertEqual(self.store.deals['one']['is_arch'], 1)
|
|
|
|
def test_nonmatching_and_other_stock_are_incremental(self):
|
|
self.snapshot(100)
|
|
self.store.sync_deals([self.deal('one', 40), self.deal('other', 60, code='600001.SH')])
|
|
self.store.archiving()
|
|
self.assertEqual(self.store.state[self.code]['base_qty'], 140)
|
|
self.assertEqual(self.store.state['600001.SH']['base_qty'], 60)
|
|
|
|
def test_matching_sell_still_liquidates(self):
|
|
self.snapshot(100)
|
|
self.store.sync_deals([self.deal('sell', 100, FLAG_SELL)])
|
|
self.store.archiving()
|
|
self.assertNotIn(self.code, self.store.state)
|
|
self.assertEqual(self.store.deals['sell']['is_arch'], 1)
|
|
|
|
def test_mixed_batch_with_matching_gross_volume_is_not_skipped(self):
|
|
self.snapshot(100)
|
|
self.store.sync_deals([self.deal('buy', 40), self.deal('sell', 60, FLAG_SELL)])
|
|
self.store.archiving()
|
|
self.assertEqual(self.store.state[self.code]['base_qty'], 80)
|
|
self.assertTrue(all(d['is_arch'] == 1 for d in self.store.deals.values()))
|
|
|
|
def test_failed_mark_rolls_back_entire_stock_and_retries(self):
|
|
self.snapshot(100)
|
|
saved = dict(self.store.state[self.code])
|
|
self.store.sync_deals([self.deal('one', 40), self.deal('two', 60)])
|
|
with closing(self.store._connect()) as db, db:
|
|
db.execute("""CREATE TRIGGER fail_mark BEFORE UPDATE OF is_arch ON deals
|
|
WHEN OLD.order_sys_id='two'
|
|
BEGIN SELECT RAISE(ABORT, 'test failure'); END""")
|
|
with self.assertLogs(level='WARNING'):
|
|
self.store.archiving()
|
|
self.assertEqual(self.store.state[self.code], saved)
|
|
self.assertTrue(all(d['is_arch'] == 0 for d in self.store.deals.values()))
|
|
with closing(self.store._connect()) as db, db:
|
|
db.execute('DROP TRIGGER fail_mark')
|
|
self.store.archiving()
|
|
self.assertEqual(self.store.state[self.code], saved)
|
|
self.assertTrue(all(d['is_arch'] == 1 for d in self.store.deals.values()))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|