Files
big-qmt/py-client/tests/test_zt_trading.py
2026-09-10 12:50:40 +08:00

137 lines
6.9 KiB
Python

import tempfile
import unittest
from datetime import datetime
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import Mock, patch
from config import AccountConfig
from libs.grid_take_profit import GridState
from libs.state import State
from sdk import Assets, DealItem, PositionItem, Tick
from strategy.zt import boot
from strategy.zt.open import open_signal
from strategy.zt.positions import manage_positions, t_rounds
class ZTTradingTests(unittest.TestCase):
def setUp(self):
tmp = tempfile.TemporaryDirectory()
self.addCleanup(tmp.cleanup)
self.store = State(Path(tmp.name) / 'state.db')
self.code = '600000.SH'
self.cfg = AccountConfig(account_id='test', buy_value=2000, zt_sell_ratio=0.5)
self.run = SimpleNamespace(account_cfg=self.cfg, orders=Mock(), client=Mock(),
profit_tracker=Mock(), add_watch=Mock(), open_watch=Mock())
self.run.orders.busy.return_value = False
self.run.orders.new_order_id.side_effect = lambda kind: f'zt-{kind}-order'
self.run.profit_tracker.observe.return_value.state = GridState.RETREAT
self.run.add_watch.triggered.return_value = True
self.run.open_watch.triggered.return_value = True
self.position = PositionItem(stock_code=self.code, volume=200, can_use_volume=200, open_price=10)
boot.sync_account_state(self.store, [self.position], [], initialize=True)
def fill(self, kind, order, qty, price=10, date='2026-09-09'):
return DealItem(stock_code=self.code, order_sys_id=order, remark=f'zt-{kind}-{order}|zt',
offset_flag=24 if kind == 't-sell' else 23,
volume=qty, price=price, trade_amount=qty * price,
trade_date=date, trade_time='100000')
def manage(self, price=11, available=10000, positions=None, force=False, today='2026-09-09'):
return manage_positions(self.run, self.store, {self.code: Tick(last_price=price)},
[self.position] if positions is None else positions,
t_rounds(self.store), available, today, force)
def test_sell_only_available_shares_and_no_loss_sell(self):
self.position.can_use_volume = 0
self.manage()
self.run.orders.place.assert_not_called()
self.position.can_use_volume = 100
self.manage(price=9)
self.run.orders.place.assert_not_called()
self.manage()
request = self.run.orders.place.call_args.args[1]
self.assertEqual((request.op, request.volume), (24, 100))
def test_full_sale_restart_and_force_buyback_without_price_or_market_gate(self):
sell = self.fill('t-sell', 's1', 200, price=11)
boot.sync_account_state(self.store, [], [sell])
self.store = State(self.store.path)
self.cfg.zt_max_price = 10
self.run.add_watch.triggered.return_value = False
remaining = self.manage(price=12, positions=[], force=True)
request = self.run.orders.place.call_args.args[1]
self.assertEqual((request.op, request.volume), (23, 200))
self.assertAlmostEqual(remaining, 10000 - 12 * 200 * 1.01)
def test_partial_fills_once_and_completed_round_blocks_same_day_sale(self):
deals = [self.fill('t-sell', 's1', 40, 11), self.fill('t-sell', 's2', 60, 12)]
self.position.volume = 100
boot.sync_account_state(self.store, [self.position], deals + deals)
item = t_rounds(self.store)[self.code]
self.assertEqual(item['sold'], 100)
self.assertEqual(item['amount'], 1160)
self.manage(price=10)
self.assertEqual(self.run.orders.place.call_args.args[1].volume, 100)
deals.append(self.fill('t-buy', 'b1', 100))
self.position.volume = 200
boot.sync_account_state(self.store, [self.position], deals)
self.run.orders.place.reset_mock()
self.manage(price=11)
self.run.orders.place.assert_not_called()
self.manage(price=11, today='2026-09-10')
self.assertEqual(self.run.orders.place.call_args.args[1].op, 24)
def test_cross_day_debt_and_insufficient_cash(self):
boot.sync_account_state(self.store, [], [self.fill('t-sell', 's1', 200, date='2026-09-08')])
self.manage(positions=[], available=100, force=True)
self.run.orders.place.assert_not_called()
self.manage(positions=[], force=True)
self.assertEqual(self.run.orders.place.call_args.args[1].volume, 200)
def test_delayed_snapshot_does_not_delete_or_recreate_holdings(self):
boot.sync_account_state(self.store, [], [])
self.assertEqual(self.store.state[self.code]['base_qty'], 200)
sell = self.fill('t-sell', 's1', 200)
boot.sync_account_state(self.store, [self.position], [sell])
self.assertNotIn(self.code, self.store.state)
self.manage()
self.run.orders.place.assert_not_called()
def test_base_fills_stay_in_base_bucket(self):
self.store.sync_state([])
deals = [self.fill('base', 'b1', 100), self.fill('base', 'b2', 100, 12)]
boot.sync_account_state(self.store, [self.position], deals)
row = self.store.state[self.code]
self.assertEqual((row['base_qty'], row['base_price'], row['added_qty']), (200, 11, 0))
def test_run_once_queries_sold_out_code_and_never_opens_with_debt(self):
sell = self.fill('t-sell', 's1', 200, 11)
self.run.client.deals.return_value = [sell]
self.run.client.portfolio.return_value = SimpleNamespace(assets=Assets(10000, 10000), positions={}, orders=[])
self.run.client.full_tick.return_value = {self.code: Tick(last_price=12)}
with patch.object(boot, 'datetime') as clock, patch.object(boot, 'collector_push'), \
patch.object(boot, 'open_signal') as opened, patch.object(boot, 'market_allow_open') as market:
clock.now.return_value = datetime(2026, 9, 9, 14, 50)
boot.RunOnce(self.run, self.store, [])
self.run.client.full_tick.assert_called_once_with([self.code])
opened.assert_not_called()
market.assert_not_called()
self.assertEqual(self.run.orders.place.call_args.args[1].op, 23)
def test_open_budget_includes_buffer_and_star_minimum(self):
with patch('strategy.zt.open.datetime') as clock:
clock.now.return_value = datetime(2026, 9, 9, 10)
remaining = open_signal(self.run, {self.code: Tick(last_price=10)},
[SimpleNamespace(code=self.code)], 2000)
self.assertEqual(self.run.orders.place.call_args.args[1].volume, 100)
self.assertEqual(remaining, 990)
self.run.orders.place.reset_mock()
open_signal(self.run, {'688001.SH': Tick(last_price=10)},
[SimpleNamespace(code='688001.SH')], 2000)
self.run.orders.place.assert_not_called()
if __name__ == '__main__':
unittest.main()