big-qmt/py-client/tests/test_deal_model.py

import ast
import tempfile
import unittest
from dataclasses import asdict, fields
from datetime import datetime
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import Mock
from libs.order import OrderBook as ActiveOrders
from libs.state import State
from sdk.models import Assets, DealItem, OrderItem, PositionItem
from sdk.portfolio import PortfolioMixin


class ApiModelTests(unittest.TestCase):
    def setUp(self):
        source = Path(__file__).resolve().parents[2] / 'api' / 'qmt_rest_new.py'
        names = {'format_assets', 'format_holding', 'format_orders', 'format_deals'}
        nodes = [n for n in ast.parse(source.read_text(encoding='utf-8')).body
                 if isinstance(n, ast.FunctionDef) and n.name in names]
        ns = {'HTTPError': RuntimeError}
        exec(compile(ast.Module(body=nodes, type_ignores=[]), str(source), 'exec'), ns)
        attrs = {n.attr: '' if n.attr.startswith('m_str') else 0
                 for node in nodes for n in ast.walk(node)
                 if isinstance(n, ast.Attribute) and n.attr.startswith('m_')}
        attrs.update(m_strInstrumentID='600000', m_strExchangeID='SH',
                     m_strOrderSysID='sys1', m_strRemark='trend-BUY-1|trend',
                     m_nOffsetFlag=23, m_nOrderStatus=56, m_nVolume=100,
                     m_nVolumeTraded=100, m_nVolumeTotalOriginal=100,
                     m_dPrice=10.0, m_dTradeAmount=1000.0, m_dBalance=2000.0,
                     m_dAvailable=1000.0, m_strInsertDate='20260907',
                     m_strInsertTime='100000', m_strTradeDate='20260907', m_strTradeTime='100000')
        obj = SimpleNamespace(**attrs)
        self.assets = ns['format_assets']([obj])
        self.positions = ns['format_holding']([obj])
        self.orders = ns['format_orders']([obj])
        self.deals = ns['format_deals']([obj])
        self.client = PortfolioMixin()
        self.client._get_json = {
            '/api/portfolio/assets': self.assets, '/api/portfolio/positions': self.positions,
            '/api/portfolio/order': self.orders, '/api/portfolio/deal': self.deals,
            '/api/portfolio': {'assets': self.assets, 'positions': self.positions, 'orders': self.orders},
        }.__getitem__

    def test_models_exactly_match_api_keys_and_values(self):
        for model, row in ((Assets, self.assets), (PositionItem, self.positions['600000.SH']),
                           (OrderItem, self.orders[0]), (DealItem, self.deals[0])):
            self.assertEqual({field.name for field in fields(model)}, set(row))
            self.assertEqual(asdict(model(**row)), row)

    def test_all_endpoints(self):
        self.assertEqual(asdict(self.client.assets()), self.assets)
        codes, positions = self.client.positions()
        self.assertEqual(codes, ['600000.SH'])
        self.assertEqual(asdict(positions[0]), self.positions[codes[0]])
        self.assertEqual(asdict(self.client.orders()[0]), self.orders[0])
        self.assertEqual(asdict(self.client.deals()[0]), self.deals[0])
        portfolio = self.client.portfolio()
        self.assertEqual(asdict(portfolio.positions[codes[0]]), self.positions[codes[0]])
        self.assertEqual(asdict(portfolio.orders[0]), self.orders[0])

    def test_derived_properties_and_order_cache(self):
        order = self.client.orders()[0]
        self.assertEqual(order.side, 'BUY')
        self.assertEqual(order.local_order_id, 'trend-BUY-1')
        self.assertEqual(order.created_at, datetime(2026, 9, 7, 10))
        order.order_status = 50
        order.insert_date = '20000101'
        client = Mock()
        book = ActiveOrders('trend')
        book.refresh(client, [order])
        client.cancel_by_id.assert_called_once_with('sys1')
        self.assertTrue(book.busy('600000.SH', 'BUY'))

    def test_storage_and_price_fallback(self):
        deal = self.client.deals()[0]
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / 'state.db'
            book = State(path)
            book.sync_deals([deal])
            loaded = State(path).deals['sys1']
            self.assertEqual(loaded['volume'], deal.volume)
            self.assertEqual(loaded['trade_date'], '2026-09-07')
            deal.order_sys_id = 'sys2'
            deal.trade_amount = 0
            book.sync_deals([deal, deal])
            self.assertEqual(book.deals['sys2']['trade_amount'], 1000)
            self.assertEqual(deal.trade_amount, 0)
            deal.order_sys_id = 'sys3'
            deal.price = 0
            with self.assertRaises(ValueError):
                book.sync_deals([deal])
            self.assertEqual(set(State(path).deals), {'sys1', 'sys2'})


if __name__ == '__main__':
    unittest.main()