From 16c237d5a0a719f7f2e2055a58077154458a93d3 Mon Sep 17 00:00:00 2001 From: minionszyw Date: Sat, 11 Apr 2026 21:47:08 +0800 Subject: [PATCH] Harden transport and decode paths --- README.md | 22 ++-- src/xmtdx/__init__.py | 2 +- src/xmtdx/_binary.py | 48 +++++++++ src/xmtdx/client.py | 74 +++++++++++--- src/xmtdx/codec/datetime_.py | 9 +- src/xmtdx/codec/frame.py | 29 +++++- src/xmtdx/codec/price.py | 28 ++++-- src/xmtdx/codec/volume.py | 5 +- src/xmtdx/commands/company_info.py | 18 ++-- src/xmtdx/commands/finance_info.py | 14 ++- src/xmtdx/commands/minute_time.py | 3 +- src/xmtdx/commands/security_bars.py | 3 +- src/xmtdx/commands/security_count.py | 5 +- src/xmtdx/commands/security_list.py | 6 +- src/xmtdx/commands/security_quotes.py | 27 +++-- src/xmtdx/commands/transaction.py | 5 +- src/xmtdx/commands/xdxr_info.py | 24 +++-- src/xmtdx/models/finance.py | 7 +- src/xmtdx/models/quote.py | 30 ++++-- src/xmtdx/models/timeseries.py | 2 +- src/xmtdx/py.typed | 1 + src/xmtdx/transport/async_.py | 74 +++++++++----- tests/integration/test_live_client.py | 36 +++++++ tests/unit/test_async_transport.py | 139 ++++++++++++++++++++++++++ tests/unit/test_codec_price.py | 1 - tests/unit/test_commands_offline.py | 4 +- tests/unit/test_decode_errors.py | 51 ++++++++++ 27 files changed, 547 insertions(+), 120 deletions(-) create mode 100644 src/xmtdx/_binary.py create mode 100644 src/xmtdx/py.typed create mode 100644 tests/integration/test_live_client.py create mode 100644 tests/unit/test_async_transport.py create mode 100644 tests/unit/test_decode_errors.py diff --git a/README.md b/README.md index ef18f16..b3863c7 100644 --- a/README.md +++ b/README.md @@ -8,11 +8,11 @@ pytdx / mootdx 年久失修：多处已知解析 bug、Python 2 包袱、无类 - **零依赖**：纯标准库，Python ≥ 3.10 - **同步 + asyncio 双接口**：`TdxClient` / `AsyncTdxClient`，commands 层不含任何 IO -- **完整类型注解**：strict mypy 通过 -- **高可用传输**：`ping_all()` 并发测速，`from_best_host()` 自动优选，断线自动重连 +- **完整类型注解**：strict `mypy` + `ruff` 通过 +- **高可用传输**：同步/异步均支持 `ping_all()`、`from_best_host()`、断线自动重连 - **修复 pytdx 已知 bug**（见下文） - **保留原始字节**：每条数据记录含 `_raw: bytes`，未知字段以 `unknown_N` 命名而非丢弃 -- **离线 fixture 测试**：13 个无网络 pytest，录制真实服务器响应验证解析正确性 +- **离线 + 本地传输回归测试**：覆盖解析、异步并发串行化、超时、自动重连与坏包处理 ## 安装 @@ -76,7 +76,8 @@ asyncio.run(main()) | `get_company_info_category(market, code)` | 公司信息文件目录 | | `get_company_info_content(market, code, filename, offset, length)` | 公司信息文本 | -`AsyncTdxClient` 提供完全相同的方法，均为 `async def`。 +`AsyncTdxClient` 提供与同步版对应的查询方法与高可用入口，均为 `async def`。 +单个 `AsyncTdxClient` 仅维护一条 TCP 连接；并发调用会在连接内串行执行。 ### KlineCategory @@ -98,6 +99,9 @@ for host, ms in results: # 自动选最优服务器 with TdxClient.from_best_host(ping_timeout=5.0) as c: ... + +# asyncio 版本同样支持 +client = AsyncTdxClient.from_best_host(ping_timeout=5.0) ``` 内置服务器列表（`KNOWN_HOSTS`）： @@ -170,6 +174,12 @@ _raw 流通股本、总股本、各省份/行业代码、资产负债表及利润表主要科目（30 个 float 字段）。 +### CompanyInfoCategory（公司信息目录） + +``` +name filename start length +``` + ## 修复的 pytdx Bug | # | 位置 | 问题 | 修复 | @@ -201,8 +211,8 @@ commands 层不依赖 transport，可独立单测。transport 层负责 TCP、 # 单元测试（无需网络） python -m pytest tests/unit/ -# 集成测试（需要网络） -python -m pytest tests/integration/ +# 集成测试（需要网络，默认跳过） +XMTDX_LIVE=1 python -m pytest tests/integration/ # 未知字段探测脚本 python scripts/probe_unknowns.py diff --git a/src/xmtdx/__init__.py b/src/xmtdx/__init__.py index d34dbbd..0f0ace1 100644 --- a/src/xmtdx/__init__.py +++ b/src/xmtdx/__init__.py @@ -22,7 +22,6 @@ asyncio 版本:: from .client import AsyncTdxClient, TdxClient from .exceptions import TdxCommandError, TdxConnectionError, TdxDecodeError, TdxError -from .transport.sync import KNOWN_HOSTS, ping_all from .models import ( XDXR_CATEGORY_NAMES, CompanyInfoCategory, @@ -36,6 +35,7 @@ from .models import ( TransactionRecord, XdxrRecord, ) +from .transport.sync import KNOWN_HOSTS, ping_all __all__ = [ # 客户端 diff --git a/src/xmtdx/_binary.py b/src/xmtdx/_binary.py new file mode 100644 index 0000000..c8b2801 --- /dev/null +++ b/src/xmtdx/_binary.py @@ -0,0 +1,48 @@ +"""二进制解析辅助函数。""" + +import struct +from typing import Any + +from .exceptions import TdxDecodeError + + +def require_bytes( + data: bytes | bytearray, + pos: int, + size: int, + context: str, +) -> None: + """确保从 pos 起至少还能读取 size 字节。""" + if pos < 0: + raise TdxDecodeError(f"{context}: 非法偏移 {pos}") + end = pos + size + if end > len(data): + remaining = max(len(data) - pos, 0) + raise TdxDecodeError( + f"{context}: 数据不足，需要 {size} 字节，偏移 {pos}，实际剩余 {remaining} 字节" + ) + + +def unpack_from( + fmt: str, + data: bytes | bytearray, + pos: int, + context: str, +) -> tuple[Any, ...]: + """带边界检查的 struct.unpack_from。""" + require_bytes(data, pos, struct.calcsize(fmt), context) + try: + return struct.unpack_from(fmt, data, pos) + except struct.error as e: # pragma: no cover - require_bytes 已覆盖大部分路径 + raise TdxDecodeError(f"{context}: 解析失败: {e}") from e + + +def slice_bytes( + data: bytes | bytearray, + pos: int, + size: int, + context: str, +) -> bytes: + """带边界检查的切片读取。""" + require_bytes(data, pos, size, context) + return bytes(data[pos : pos + size]) diff --git a/src/xmtdx/client.py b/src/xmtdx/client.py index b10fb28..6e996c7 100644 --- a/src/xmtdx/client.py +++ b/src/xmtdx/client.py @@ -1,5 +1,6 @@ """高层行情 API：TdxClient（同步）和 AsyncTdxClient（asyncio）。""" +import asyncio from types import TracebackType from typing import TypeVar @@ -243,6 +244,9 @@ class AsyncTdxClient: async with AsyncTdxClient("180.153.18.170") as c: bars = await c.get_security_bars(Market.SH, "600000", KlineCategory.DAY, 0, 100) + + 注意： + 单个 AsyncTdxClient 仅维护一条 TCP 连接；并发调用会在连接内串行执行。 """ def __init__( @@ -250,8 +254,37 @@ class AsyncTdxClient: host: str = KNOWN_HOSTS[0], port: int = _DEFAULT_PORT, timeout: float = 15.0, + auto_reconnect: bool = True, ) -> None: + self._host = host + self._port = port + self._timeout = timeout + self._auto_reconnect = auto_reconnect self._conn = AsyncTdxConnection(host, port, timeout) + self._execute_lock = asyncio.Lock() + + @classmethod + def from_best_host( + cls, + hosts: list[str] = KNOWN_HOSTS, + port: int = _DEFAULT_PORT, + timeout: float = 15.0, + ping_timeout: float = 5.0, + auto_reconnect: bool = True, + ) -> "AsyncTdxClient": + """测量 hosts 中所有服务器延迟，选最低延迟的建立连接。""" + ranked = ping_all(hosts, port, ping_timeout) + best = ranked[0][0] if ranked else hosts[0] + return cls(best, port, timeout, auto_reconnect) + + @staticmethod + def ping_all( + hosts: list[str] = KNOWN_HOSTS, + port: int = _DEFAULT_PORT, + timeout: float = 5.0, + ) -> list[tuple[str, float]]: + """测量多台服务器延迟，返回按延迟排序的 (host, seconds) 列表。""" + return ping_all(hosts, port, timeout) async def connect(self) -> None: await self._conn.connect() @@ -271,16 +304,29 @@ class AsyncTdxClient: ) -> None: await self.close() + async def _execute(self, cmd: "BaseCommand[_T]") -> _T: + """执行命令；断线时尝试重连一次再重试（若 auto_reconnect=True）。""" + async with self._execute_lock: + try: + return await self._conn.execute(cmd) + except TdxConnectionError: + if not self._auto_reconnect: + raise + await self._conn.close() + self._conn = AsyncTdxConnection(self._host, self._port, self._timeout) + await self._conn.connect() + return await self._conn.execute(cmd) + async def get_security_count(self, market: Market) -> int: - return await self._conn.execute(GetSecurityCountCmd(market)) + return await self._execute(GetSecurityCountCmd(market)) async def get_security_list(self, market: Market, start: int) -> list[SecurityInfo]: - return await self._conn.execute(GetSecurityListCmd(market, start)) + return await self._execute(GetSecurityListCmd(market, start)) async def get_security_quotes( self, stocks: list[tuple[Market, str]] ) -> list[SecurityQuote]: - return await self._conn.execute(GetSecurityQuotesCmd(stocks)) + return await self._execute(GetSecurityQuotesCmd(stocks)) async def get_security_bars( self, @@ -290,7 +336,9 @@ class AsyncTdxClient: start: int, count: int = 800, ) -> list[SecurityBar]: - return await self._conn.execute(GetSecurityBarsCmd(market, code, category, start, count)) + return await self._execute( + GetSecurityBarsCmd(market, code, category, start, count) + ) async def get_index_bars( self, @@ -300,42 +348,42 @@ class AsyncTdxClient: start: int, count: int = 800, ) -> list[SecurityBar]: - return await self._conn.execute(GetIndexBarsCmd(market, code, category, start, count)) + return await self._execute(GetIndexBarsCmd(market, code, category, start, count)) async def get_minute_time_data(self, market: Market, code: str) -> list[MinuteBar]: - return await self._conn.execute(GetMinuteTimeDataCmd(market, code)) + return await self._execute(GetMinuteTimeDataCmd(market, code)) async def get_history_minute_time_data( self, market: Market, code: str, date: int ) -> list[MinuteBar]: - return await self._conn.execute(GetHistoryMinuteTimeDataCmd(market, code, date)) + return await self._execute(GetHistoryMinuteTimeDataCmd(market, code, date)) async def get_transaction_data( self, market: Market, code: str, start: int, count: int = 800 ) -> list[TransactionRecord]: - return await self._conn.execute(GetTransactionDataCmd(market, code, start, count)) + return await self._execute(GetTransactionDataCmd(market, code, start, count)) async def get_history_transaction_data( self, market: Market, code: str, date: int, start: int, count: int = 800 ) -> list[TransactionRecord]: - return await self._conn.execute( + return await self._execute( GetHistoryTransactionDataCmd(market, code, date, start, count) ) async def get_xdxr_info(self, market: Market, code: str) -> list[XdxrRecord]: - return await self._conn.execute(GetXdxrInfoCmd(market, code)) + return await self._execute(GetXdxrInfoCmd(market, code)) async def get_finance_info(self, market: Market, code: str) -> FinanceInfo: - return await self._conn.execute(GetFinanceInfoCmd(market, code)) + return await self._execute(GetFinanceInfoCmd(market, code)) async def get_company_info_category( self, market: Market, code: str ) -> list[CompanyInfoCategory]: - return await self._conn.execute(GetCompanyInfoCategoryCmd(market, code)) + return await self._execute(GetCompanyInfoCategoryCmd(market, code)) async def get_company_info_content( self, market: Market, code: str, filename: str, offset: int, length: int ) -> str: - return await self._conn.execute( + return await self._execute( GetCompanyInfoContentCmd(market, code, filename, offset, length) ) diff --git a/src/xmtdx/codec/datetime_.py b/src/xmtdx/codec/datetime_.py index 0ea365c..78ad1fa 100644 --- a/src/xmtdx/codec/datetime_.py +++ b/src/xmtdx/codec/datetime_.py @@ -7,7 +7,8 @@ 日线及以上（其余 category）：4 字节 YYYYMMDD 整数 """ -import struct + +from .._binary import unpack_from def get_datetime_minute( @@ -18,7 +19,7 @@ def get_datetime_minute( Returns: (year, month, day, hour, minute, new_pos) """ - zipday, tminutes = struct.unpack_from("> 11) + 2004 month = (zipday % 2048) // 100 day = (zipday % 2048) % 100 @@ -35,7 +36,7 @@ def get_datetime_day( Returns: (year, month, day, new_pos) """ - (zipday,) = struct.unpack_from(" tuple[int, int, int]: Returns: (hour, minute, new_pos) """ - (tminutes,) = struct.unpack_from(" FrameHeader: """解析 16 字节响应帧头。""" - u0, u1, u2, zipsize, unzipsize = struct.unpack_from(_HEADER_FMT, buf) + u0, u1, u2, zipsize, unzipsize = unpack_from( + _HEADER_FMT, + buf, + 0, + "frame header", + ) return FrameHeader(u0, u1, u2, zipsize, unzipsize) @@ -37,6 +44,20 @@ def decompress_body(header: FrameHeader, raw_body: bytes) -> bytes: zipsize == unzipsize 时直接返回原始字节；否则 zlib 解压。 """ + if len(raw_body) != header.zipsize: + raise TdxDecodeError( + f"frame body 长度不符: header={header.zipsize}, actual={len(raw_body)}" + ) if header.zipsize == header.unzipsize: - return raw_body - return zlib.decompress(raw_body) + body = raw_body + else: + try: + body = zlib.decompress(raw_body) + except zlib.error as e: + raise TdxDecodeError(f"frame body zlib 解压失败: {e}") from e + + if len(body) != header.unzipsize: + raise TdxDecodeError( + f"frame body 解压长度不符: header={header.unzipsize}, actual={len(body)}" + ) + return body diff --git a/src/xmtdx/codec/price.py b/src/xmtdx/codec/price.py index 2faa2b2..a312943 100644 --- a/src/xmtdx/codec/price.py +++ b/src/xmtdx/codec/price.py @@ -8,6 +8,8 @@ 典型用途：价格差分、成交量差分、买卖档位数量。 """ +from ..exceptions import TdxDecodeError + def get_price(data: bytes | bytearray, pos: int) -> tuple[int, int]: """解码一个变长有符号整数。 @@ -16,18 +18,22 @@ def get_price(data: bytes | bytearray, pos: int) -> tuple[int, int]: (value, new_pos) """ bit_shift = 6 - b = data[pos] - value = b & 0x3F - negative = bool(b & 0x40) + start = pos + try: + b = data[pos] + value = b & 0x3F + negative = bool(b & 0x40) - if b & 0x80: - while True: - pos += 1 - b = data[pos] - value |= (b & 0x7F) << bit_shift - bit_shift += 7 - if not (b & 0x80): - break + if b & 0x80: + while True: + pos += 1 + b = data[pos] + value |= (b & 0x7F) << bit_shift + bit_shift += 7 + if not (b & 0x80): + break + except IndexError as e: + raise TdxDecodeError(f"price varint 截断: offset={start}") from e pos += 1 return (-value if negative else value), pos diff --git a/src/xmtdx/codec/volume.py b/src/xmtdx/codec/volume.py index 749a356..c7b1b26 100644 --- a/src/xmtdx/codec/volume.py +++ b/src/xmtdx/codec/volume.py @@ -9,7 +9,8 @@ 警告：此函数专为成交量设计，不可用于价格字段（pytdx Bug #3）。 """ -import struct + +from .._binary import unpack_from def get_volume(data: bytes | bytearray, pos: int) -> tuple[float, int]: @@ -18,7 +19,7 @@ def get_volume(data: bytes | bytearray, pos: int) -> tuple[float, int]: Returns: (volume_float, new_pos) """ - (ivol,) = struct.unpack_from(" list[CompanyInfoCategory]: if len(body) < 2: - return [] - (num,) = struct.unpack_from(" len(body): - break - name_b, filename_b, start, length = struct.unpack_from("<64s80sII", body, pos) + raw = slice_bytes(body, pos, _RECORD_SIZE, "company_info_category record") + name_b, filename_b, start, length = struct.unpack("<64s80sII", raw) pos += _RECORD_SIZE def _decode(b: bytes) -> str: @@ -39,6 +40,7 @@ class GetCompanyInfoCategoryCmd(BaseCommand[list[CompanyInfoCategory]]): return raw.decode("gbk", errors="replace") results.append(CompanyInfoCategory( + name=_decode(name_b), filename=_decode(filename_b), start=start, length=length, @@ -70,7 +72,7 @@ class GetCompanyInfoContentCmd(BaseCommand[str]): def parse_response(self, body: bytes) -> str: # 前12字节：10字节未知 + 2字节长度 if len(body) < 12: - return "" - _, length = struct.unpack_from("<10sH", body, 0) - content = body[12 : 12 + length] + raise TdxDecodeError("company_info_content body 过短") + _, length = unpack_from("<10sH", body, 0, "company_info_content header") + content = slice_bytes(body, 12, length, "company_info_content body") return content.decode("gbk", errors="replace") diff --git a/src/xmtdx/commands/finance_info.py b/src/xmtdx/commands/finance_info.py index 030101e..be085ae 100644 --- a/src/xmtdx/commands/finance_info.py +++ b/src/xmtdx/commands/finance_info.py @@ -2,6 +2,8 @@ import struct +from .._binary import slice_bytes, unpack_from +from ..exceptions import TdxDecodeError from ..models.enums import Market from ..models.finance import FinanceInfo from .base import BaseCommand @@ -24,10 +26,10 @@ class GetFinanceInfoCmd(BaseCommand[FinanceInfo]): def parse_response(self, body: bytes) -> FinanceInfo: pos = 2 # 跳过前2字节（记录数） - market_b, code_b = struct.unpack_from(" list[MinuteBar]: - (num,) = struct.unpack_from(" list[SecurityBar]: - (ret_count,) = struct.unpack_from(" int: - (count,) = struct.unpack_from(" list[SecurityInfo]: - (num,) = struct.unpack_from(" list[TransactionRecord]: """当日逐笔：time + price + vol + num_orders + buyorsell + unknown""" - (num,) = struct.unpack_from(" list[TransactionRecord]: def _parse_history_transaction_body(body: bytes) -> list[TransactionRecord]: """历史逐笔：num(2) + skip(4) + [time + price + vol + buyorsell + unknown]""" - (num,) = struct.unpack_from(" list[XdxrRecord]: if len(body) < 11: - return [] + raise TdxDecodeError("xdxr_info body 过短") pos = 9 # 跳过9字节（market+code+未知） - (num,) = struct.unpack_from(" len(body): - break - market_b, code_b = struct.unpack_from(" len(body): - break - - chunk = body[pos : pos + 16] + chunk = slice_bytes(body, pos, 16, "xdxr_info record body") pos += 16 + try: + market = Market(market_b) + except ValueError as e: + raise TdxDecodeError(f"xdxr_info 非法 market 值: {market_b}") from e rec = XdxrRecord( - market=Market(market_b), + market=market, code=code_b.decode("utf-8").rstrip("\x00"), year=year, month=month, diff --git a/src/xmtdx/models/finance.py b/src/xmtdx/models/finance.py index 6b6a75c..f0ab72d 100644 --- a/src/xmtdx/models/finance.py +++ b/src/xmtdx/models/finance.py @@ -122,6 +122,7 @@ class FinanceInfo: class CompanyInfoCategory: """公司信息文件目录条目""" - filename: str # 文件名（如 '600000.txt'） - start: int # 内容起始偏移 - length: int # 内容长度（字节） + name: str = "" # 目录名（如“最新提示”） + filename: str = "" # 文件名（如 '600000.txt'） + start: int = 0 # 内容起始偏移 + length: int = 0 # 内容长度（字节） diff --git a/src/xmtdx/models/quote.py b/src/xmtdx/models/quote.py index 3f66e26..1c78e9e 100644 --- a/src/xmtdx/models/quote.py +++ b/src/xmtdx/models/quote.py @@ -35,18 +35,28 @@ class SecurityQuote: active2: int # 买盘五档 - bid1: float; bid_vol1: float - bid2: float; bid_vol2: float - bid3: float; bid_vol3: float - bid4: float; bid_vol4: float - bid5: float; bid_vol5: float + bid1: float + bid_vol1: float + bid2: float + bid_vol2: float + bid3: float + bid_vol3: float + bid4: float + bid_vol4: float + bid5: float + bid_vol5: float # 卖盘五档 - ask1: float; ask_vol1: float - ask2: float; ask_vol2: float - ask3: float; ask_vol3: float - ask4: float; ask_vol4: float - ask5: float; ask_vol5: float + ask1: float + ask_vol1: float + ask2: float + ask_vol2: float + ask3: float + ask_vol3: float + ask4: float + ask_vol4: float + ask5: float + ask_vol5: float # 已确认含义 rise_speed: float # 涨速（原 reversed_bytes9 / 100） diff --git a/src/xmtdx/models/timeseries.py b/src/xmtdx/models/timeseries.py index 985a006..4b1db4e 100644 --- a/src/xmtdx/models/timeseries.py +++ b/src/xmtdx/models/timeseries.py @@ -31,7 +31,7 @@ class TransactionRecord: minute: int price: float vol: int - buyorsell: int # 0=卖, 1=买, 2=中性/撮合 + buyorsell: int # 0=买, 1=卖, 2=中性/撮合, 8=集合竞价 # pytdx 中被丢弃的字段 unknown_last: int = field(default=0, repr=False) diff --git a/src/xmtdx/py.typed b/src/xmtdx/py.typed new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/src/xmtdx/py.typed @@ -0,0 +1 @@ + diff --git a/src/xmtdx/transport/async_.py b/src/xmtdx/transport/async_.py index 80fc22a..bdfcde9 100644 --- a/src/xmtdx/transport/async_.py +++ b/src/xmtdx/transport/async_.py @@ -38,9 +38,47 @@ class AsyncTdxConnection: self.timeout = timeout self._reader: asyncio.StreamReader | None = None self._writer: asyncio.StreamWriter | None = None + # 单连接不支持请求复用；所有 IO 在连接内串行执行。 + self._io_lock = asyncio.Lock() async def connect(self) -> None: """建立 TCP 连接并完成握手。""" + async with self._io_lock: + if self._writer is not None and not self._writer.is_closing(): + return + await self._connect_unlocked() + + async def close(self) -> None: + """关闭连接。""" + async with self._io_lock: + await self._close_unlocked() + + async def execute(self, cmd: "BaseCommand[T]") -> T: + """执行一条命令（异步版本）。 + + 同一连接上的并发调用会在此处串行化，避免 StreamReader 并发读取冲突。 + """ + async with self._io_lock: + if self._writer is None or self._reader is None: + raise TdxConnectionError("未连接，请先调用 connect()") + request = cmd.build_request() + try: + self._writer.write(request) + await asyncio.wait_for(self._writer.drain(), timeout=self.timeout) + header_buf = await self._recv_exact(HEADER_SIZE) + header = parse_header(header_buf) + raw_body = await self._recv_exact(header.zipsize) + except asyncio.TimeoutError as e: + await self._close_unlocked() + raise TdxConnectionError(f"通信超时: {self.timeout}s") from e + except (OSError, asyncio.IncompleteReadError) as e: + await self._close_unlocked() + raise TdxConnectionError(f"通信错误: {e}") from e + + body = decompress_body(header, raw_body) + return cmd.parse_response(body) + + async def _connect_unlocked(self) -> None: try: reader, writer = await asyncio.wait_for( asyncio.open_connection(self.host, self.port), @@ -50,10 +88,13 @@ class AsyncTdxConnection: raise TdxConnectionError(f"无法连接 {self.host}:{self.port}: {e}") from e self._reader = reader self._writer = writer - await self._send_setup() + try: + await self._send_setup() + except Exception: + await self._close_unlocked() + raise - async def close(self) -> None: - """关闭连接。""" + async def _close_unlocked(self) -> None: if self._writer is not None: try: self._writer.close() @@ -63,22 +104,6 @@ class AsyncTdxConnection: self._reader = None self._writer = None - async def execute(self, cmd: "BaseCommand[T]") -> T: - """执行一条命令（异步版本）。""" - if self._writer is None or self._reader is None: - raise TdxConnectionError("未连接，请先调用 connect()") - request = cmd.build_request() - try: - self._writer.write(request) - await self._writer.drain() - header_buf = await self._recv_exact(HEADER_SIZE) - header = parse_header(header_buf) - raw_body = await self._recv_exact(header.zipsize) - except (OSError, asyncio.IncompleteReadError) as e: - raise TdxConnectionError(f"通信错误: {e}") from e - body = decompress_body(header, raw_body) - return cmd.parse_response(body) - # ------------------------------------------------------------------ # # context manager # ------------------------------------------------------------------ # @@ -105,11 +130,9 @@ class AsyncTdxConnection: assert self._reader is not None for cmd_bytes in SETUP_COMMANDS: self._writer.write(cmd_bytes) - await self._writer.drain() + await asyncio.wait_for(self._writer.drain(), timeout=self.timeout) try: - hdr_buf = await asyncio.wait_for( - self._recv_exact(HEADER_SIZE), timeout=5.0 - ) + hdr_buf = await self._recv_exact(HEADER_SIZE) hdr = parse_header(hdr_buf) if hdr.zipsize > 0: await self._recv_exact(hdr.zipsize) @@ -119,5 +142,8 @@ class AsyncTdxConnection: async def _recv_exact(self, n: int) -> bytes: """读满 n 字节。""" assert self._reader is not None - data = await self._reader.readexactly(n) + data = await asyncio.wait_for( + self._reader.readexactly(n), + timeout=self.timeout, + ) return data diff --git a/tests/integration/test_live_client.py b/tests/integration/test_live_client.py new file mode 100644 index 0000000..a0a9a46 --- /dev/null +++ b/tests/integration/test_live_client.py @@ -0,0 +1,36 @@ +"""真实通达信服务器 smoke test。 + +默认跳过；设置 XMTDX_LIVE=1 后执行。 +""" + +from __future__ import annotations + +import asyncio +import os + +import pytest + +from xmtdx import AsyncTdxClient, Market, TdxClient + +_LIVE_ENABLED = os.getenv("XMTDX_LIVE") == "1" +_LIVE_HOST = os.getenv("XMTDX_HOST", "180.153.18.170") + +pytestmark = pytest.mark.skipif( + not _LIVE_ENABLED, + reason="set XMTDX_LIVE=1 to run live integration tests", +) + + +def test_sync_live_smoke() -> None: + with TdxClient(_LIVE_HOST, timeout=5.0) as client: + assert client.get_security_count(Market.SH) > 0 + assert client.get_security_count(Market.SZ) > 0 + + +def test_async_live_smoke() -> None: + async def main() -> None: + async with AsyncTdxClient(_LIVE_HOST, timeout=5.0) as client: + assert await client.get_security_count(Market.SH) > 0 + assert await client.get_security_count(Market.SZ) > 0 + + asyncio.run(main()) diff --git a/tests/unit/test_async_transport.py b/tests/unit/test_async_transport.py new file mode 100644 index 0000000..e53994e --- /dev/null +++ b/tests/unit/test_async_transport.py @@ -0,0 +1,139 @@ +"""异步 transport 回归测试。""" + +from __future__ import annotations + +import asyncio +import struct +import time + +from xmtdx import AsyncTdxClient, Market +from xmtdx.commands.security_count import GetSecurityCountCmd +from xmtdx.commands.setup import SETUP_COMMANDS +from xmtdx.exceptions import TdxConnectionError + + +def _pack_frame(body: bytes) -> bytes: + return struct.pack(" None: + request_len = len(GetSecurityCountCmd(Market.SH).build_request()) + + async def handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + try: + for setup_cmd in SETUP_COMMANDS: + await reader.readexactly(len(setup_cmd)) + writer.write(_pack_frame(b"")) + await writer.drain() + + await reader.readexactly(request_len) + writer.write(_pack_frame(struct.pack(" None: + server = await asyncio.start_server(handle, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + try: + client = AsyncTdxClient("127.0.0.1", port=port, timeout=0.2) + await client.connect() + try: + sh_count, sz_count = await asyncio.gather( + client.get_security_count(Market.SH), + client.get_security_count(Market.SZ), + ) + finally: + await client.close() + finally: + server.close() + await server.wait_closed() + + assert sh_count == 5 + assert sz_count == 6 + + asyncio.run(main()) + + +def test_async_client_auto_reconnect() -> None: + request_len = len(GetSecurityCountCmd(Market.SH).build_request()) + connection_ids: list[int] = [] + + async def handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + connection_ids.append(len(connection_ids) + 1) + connection_id = connection_ids[-1] + try: + for setup_cmd in SETUP_COMMANDS: + await reader.readexactly(len(setup_cmd)) + writer.write(_pack_frame(b"")) + await writer.drain() + + await reader.readexactly(request_len) + writer.write(_pack_frame(struct.pack(" None: + server = await asyncio.start_server(handle, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + try: + client = AsyncTdxClient("127.0.0.1", port=port, timeout=0.2) + first = await client.get_security_count(Market.SH) + second = await client.get_security_count(Market.SH) + await client.close() + finally: + server.close() + await server.wait_closed() + + assert first == 11 + assert second == 12 + assert len(connection_ids) == 2 + + asyncio.run(main()) + + +def test_async_client_request_timeout() -> None: + request_len = len(GetSecurityCountCmd(Market.SH).build_request()) + + async def handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + try: + for setup_cmd in SETUP_COMMANDS: + await reader.readexactly(len(setup_cmd)) + writer.write(_pack_frame(b"")) + await writer.drain() + + await reader.readexactly(request_len) + await asyncio.sleep(1.0) + finally: + writer.close() + await writer.wait_closed() + + async def main() -> None: + server = await asyncio.start_server(handle, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + try: + client = AsyncTdxClient("127.0.0.1", port=port, timeout=0.05) + await client.connect() + t0 = time.monotonic() + try: + await client.get_security_count(Market.SH) + except TdxConnectionError as exc: + elapsed = time.monotonic() - t0 + assert "超时" in str(exc) or "timed out" in str(exc) + assert elapsed < 0.3 + else: # pragma: no cover - 防御性断言 + raise AssertionError("expected timeout") + finally: + await client.close() + finally: + server.close() + await server.wait_closed() + + asyncio.run(main()) diff --git a/tests/unit/test_codec_price.py b/tests/unit/test_codec_price.py index 9a217db..176abf4 100644 --- a/tests/unit/test_codec_price.py +++ b/tests/unit/test_codec_price.py @@ -1,6 +1,5 @@ """get_price / put_price 单元测试，测试向量来自 pytdx 实际报文。""" -import pytest from xmtdx.codec.price import get_price, put_price diff --git a/tests/unit/test_commands_offline.py b/tests/unit/test_commands_offline.py index bfe4aa6..25352b7 100644 --- a/tests/unit/test_commands_offline.py +++ b/tests/unit/test_commands_offline.py @@ -7,7 +7,6 @@ fixtures/ 目录下每个 .hex 文件是一次真实服务器响应的 body（ from __future__ import annotations import pathlib -import pytest FIXTURES = pathlib.Path(__file__).parent.parent / "fixtures" @@ -75,7 +74,7 @@ def test_security_list_gbk_no_crash(): def test_security_bars_parse(): from xmtdx.commands.security_bars import GetSecurityBarsCmd - from xmtdx.models.enums import Market, KlineCategory + from xmtdx.models.enums import KlineCategory, Market body = load_hex("security_bars") cmd = GetSecurityBarsCmd(Market.SH, "600000", KlineCategory.DAY, 0, 5) @@ -305,6 +304,7 @@ def test_company_info_category_parse(): assert len(cats) == 16 c0 = cats[0] + assert c0.name == "最新提示" assert c0.filename == "600000.txt" assert c0.start == 0 assert c0.length == 11426 diff --git a/tests/unit/test_decode_errors.py b/tests/unit/test_decode_errors.py new file mode 100644 index 0000000..f10fa2f --- /dev/null +++ b/tests/unit/test_decode_errors.py @@ -0,0 +1,51 @@ +"""坏包与解码异常回归测试。""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from xmtdx.codec.frame import FrameHeader, decompress_body +from xmtdx.commands.company_info import GetCompanyInfoCategoryCmd +from xmtdx.commands.security_count import GetSecurityCountCmd +from xmtdx.commands.xdxr_info import GetXdxrInfoCmd +from xmtdx.exceptions import TdxDecodeError +from xmtdx.models.enums import Market + +FIXTURES = Path(__file__).parent.parent / "fixtures" + + +def _load_hex(name: str) -> bytes: + return bytes.fromhex((FIXTURES / f"{name}.hex").read_text().strip()) + + +def test_security_count_truncated_raises_tdxdecodeerror() -> None: + with pytest.raises(TdxDecodeError): + GetSecurityCountCmd(Market.SH).parse_response(b"") + + +def test_company_info_category_truncated_raises_tdxdecodeerror() -> None: + body = _load_hex("company_info_category") + cmd = GetCompanyInfoCategoryCmd(Market.SH, "600000") + + with pytest.raises(TdxDecodeError): + cmd.parse_response(body[:-10]) + + +def test_xdxr_info_truncated_raises_tdxdecodeerror() -> None: + body = _load_hex("xdxr_info") + cmd = GetXdxrInfoCmd(Market.SH, "600000") + + with pytest.raises(TdxDecodeError): + cmd.parse_response(body[:-10]) + + +def test_frame_bad_zlib_raises_tdxdecodeerror() -> None: + with pytest.raises(TdxDecodeError): + decompress_body(FrameHeader(0, 0, 0, 4, 8), b"xxxx") + + +def test_frame_unzipsize_mismatch_raises_tdxdecodeerror() -> None: + with pytest.raises(TdxDecodeError): + decompress_body(FrameHeader(0, 0, 0, 3, 4), b"abc")