Files
big-qmt/py-client/tests/test_ipo.py
2026-08-29 01:50:09 +08:00

141 lines
4.9 KiB
Python

from __future__ import annotations
import unittest
from datetime import datetime
from tempfile import TemporaryDirectory
from types import SimpleNamespace
from unittest.mock import patch
from strategy.ipo.boot import AutoBuyIpo
RUN_TIME = datetime(2026, 8, 28, 10, 0)
class FakeClient:
def __init__(self, candidates=None, orders=None, deals=None, fail_codes=None):
self.candidates = candidates or {}
self.orders = orders or []
self.deal_rows = deals or []
self.fail_codes = set(fail_codes or [])
self.submissions = []
self.closed = False
def __enter__(self):
return self
def __exit__(self, *_args):
self.closed = True
def trading_dates(self, *_args):
return ["20260828"]
def trade_detail_data(self, datatype):
self.assert_order_type = datatype
return self.orders
def deals(self):
return self.deal_rows
def ipo_data(self, ipo_type):
self.assert_ipo_type = ipo_type
return self.candidates
def passorder(self, **kwargs):
code = kwargs["stock"]
self.submissions.append(kwargs)
if code in self.fail_codes:
raise RuntimeError("simulated rejection")
return {"status": "success", "order_ref": f"ref-{code}"}
class AutoBuyIpoTests(unittest.TestCase):
def _configs(self, directory, enabled=True):
return (
SimpleNamespace(
qmt_base_url="http://qmt",
qmt_token="token",
qmt_data_dir=directory,
),
SimpleNamespace(account_id="account-A", enable_auto_ipo=enabled),
)
def test_disabled_does_not_create_client(self):
with TemporaryDirectory() as directory:
global_cfg, account_cfg = self._configs(directory, enabled=False)
with (
patch("strategy.ipo.boot.config.global_config", global_cfg),
patch("strategy.ipo.boot.config.account_config", account_cfg),
patch("strategy.ipo.boot.Client") as client_factory,
):
self.assertEqual(AutoBuyIpo(RUN_TIME), 0)
client_factory.assert_not_called()
def test_local_record_prevents_duplicate_after_restart(self):
candidates = {
"688001.SH": {"issuePrice": 10, "maxPurchaseNum": 1000},
}
first = FakeClient(candidates=candidates)
second = FakeClient(candidates=candidates)
with TemporaryDirectory() as directory:
global_cfg, account_cfg = self._configs(directory)
with (
patch("strategy.ipo.boot.config.global_config", global_cfg),
patch("strategy.ipo.boot.config.account_config", account_cfg),
patch("strategy.ipo.boot.Client", side_effect=[first, second]),
):
self.assertEqual(AutoBuyIpo(RUN_TIME), 1)
self.assertEqual(AutoBuyIpo(RUN_TIME), 0)
self.assertEqual(len(first.submissions), 1)
self.assertEqual(second.submissions, [])
self.assertTrue(first.closed)
self.assertTrue(second.closed)
def test_broker_order_prevents_duplicate(self):
candidates = {
"688001.SH": {"issuePrice": 10, "maxPurchaseNum": 1000},
}
client = FakeClient(
candidates=candidates,
orders=[{
"m_strInstrumentID": "688001",
"m_strInsertDate": "20260828",
"m_strRemark": "IPO_SUBSCRIBE",
}],
)
with TemporaryDirectory() as directory:
global_cfg, account_cfg = self._configs(directory)
with (
patch("strategy.ipo.boot.config.global_config", global_cfg),
patch("strategy.ipo.boot.config.account_config", account_cfg),
patch("strategy.ipo.boot.Client", return_value=client),
):
self.assertEqual(AutoBuyIpo(RUN_TIME), 0)
self.assertEqual(client.submissions, [])
def test_one_rejection_does_not_stop_other_candidates(self):
candidates = {
"688001.SH": {"issuePrice": 10, "maxPurchaseNum": 1000},
"688002.SH": {"issuePrice": 20, "maxPurchaseNum": 500},
}
client = FakeClient(candidates=candidates, fail_codes={"688001.SH"})
with TemporaryDirectory() as directory:
global_cfg, account_cfg = self._configs(directory)
with (
patch("strategy.ipo.boot.config.global_config", global_cfg),
patch("strategy.ipo.boot.config.account_config", account_cfg),
patch("strategy.ipo.boot.Client", return_value=client),
):
self.assertEqual(AutoBuyIpo(RUN_TIME), 1)
self.assertEqual(
[item["stock"] for item in client.submissions],
["688001.SH", "688002.SH"],
)
self.assertTrue(client.closed)
if __name__ == "__main__":
unittest.main()