fix zt&state.py

py-client/tests/test_state_snapshot.py

@@ -0,0 +1,92 @@
+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()