From 081f851c6ddc40fa0c73a1b7bde2e8fbf7d1e31e Mon Sep 17 00:00:00 2001
From: M <m@example.com>
Date: Tue, 14 Apr 2026 18:11:05 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20AsyncTdxClient=20=E5=A2=9E=E5=8A=A0?=
 =?UTF-8?q?=E5=BF=83=E8=B7=B3=E4=BF=9D=E6=B4=BB=E6=9C=BA=E5=88=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

1. 在 AsyncTdxClient 中增加后台心跳任务，默认每 60 秒发送一次请求。
2. 支持在连接成功后自动启动，并在连接关闭时自动停止心跳任务。
3. 增加单元测试 tests/unit/test_heartbeat.py 验证心跳循环与清理逻辑。
---
 src/xmtdx/client.py          | 40 +++++++++++++++++++++-
 tests/unit/test_heartbeat.py | 65 ++++++++++++++++++++++++++++++++++++
 2 files changed, 104 insertions(+), 1 deletion(-)
 create mode 100644 tests/unit/test_heartbeat.py

diff --git a/src/xmtdx/client.py b/src/xmtdx/client.py
index 163376c..c4f97ec 100644
--- a/src/xmtdx/client.py
+++ b/src/xmtdx/client.py
@@ -293,13 +293,16 @@ class AsyncTdxClient:
         port: int = _DEFAULT_PORT,
         timeout: float = 15.0,
         auto_reconnect: bool = True,
+        heartbeat_interval: float = 60.0,
     ) -> None:
         self._host = host
         self._port = port
         self._timeout = timeout
         self._auto_reconnect = auto_reconnect
+        self._heartbeat_interval = heartbeat_interval
         self._conn = AsyncTdxConnection(host, port, timeout)
         self._execute_lock = asyncio.Lock()
+        self._heartbeat_task: asyncio.Task[None] | None = None
 
     @classmethod
     def from_best_host(
@@ -309,11 +312,12 @@ class AsyncTdxClient:
         timeout: float = 15.0,
         ping_timeout: float = 5.0,
         auto_reconnect: bool = True,
+        heartbeat_interval: float = 60.0,
     ) -> "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)
+        return cls(best, port, timeout, auto_reconnect, heartbeat_interval)
 
     @staticmethod
     def ping_all(
@@ -326,8 +330,10 @@ class AsyncTdxClient:
 
     async def connect(self) -> None:
         await self._conn.connect()
+        self._start_heartbeat()
 
     async def close(self) -> None:
+        await self._stop_heartbeat()
         await self._conn.close()
 
     async def __aenter__(self) -> "AsyncTdxClient":
@@ -342,6 +348,38 @@ class AsyncTdxClient:
     ) -> None:
         await self.close()
 
+    def _start_heartbeat(self) -> None:
+        """启动后台心跳任务。"""
+        if self._heartbeat_interval <= 0:
+            return
+        if self._heartbeat_task is not None:
+            self._heartbeat_task.cancel()
+        self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
+
+    async def _stop_heartbeat(self) -> None:
+        """停止并清理心跳任务。"""
+        if self._heartbeat_task:
+            self._heartbeat_task.cancel()
+            try:
+                await self._heartbeat_task
+            except asyncio.CancelledError:
+                pass
+            self._heartbeat_task = None
+
+    async def _heartbeat_loop(self) -> None:
+        """心跳循环：定期发送轻量级请求保活。"""
+        while True:
+            try:
+                await asyncio.sleep(self._heartbeat_interval)
+                # 使用 get_security_count 作为心跳包
+                await self.get_security_count(Market.SH)
+            except asyncio.CancelledError:
+                break
+            except Exception:
+                # 心跳失败通常意味着连接已断开
+                # 下一次正常的业务请求或下一次心跳会通过 _execute 触发重连
+                pass
+
     async def _execute(self, cmd: "BaseCommand[_T]") -> _T:
         """执行命令；断线时尝试重连一次再重试（若 auto_reconnect=True）。"""
         async with self._execute_lock:
diff --git a/tests/unit/test_heartbeat.py b/tests/unit/test_heartbeat.py
new file mode 100644
index 0000000..4e99241
--- /dev/null
+++ b/tests/unit/test_heartbeat.py
@@ -0,0 +1,65 @@
+"""心跳机制单元测试。"""
+
+import asyncio
+import pytest
+from unittest.mock import patch, MagicMock, AsyncMock
+from xmtdx import AsyncTdxClient, Market
+
+
+@pytest.mark.asyncio
+async def test_heartbeat_sends_periodically():
+    # 模拟连接和执行
+    with patch("xmtdx.client.AsyncTdxConnection") as mock_conn_cls:
+        mock_conn = mock_conn_cls.return_value
+        mock_conn.connect = AsyncMock()
+        mock_conn.close = AsyncMock()
+        
+        # 记录调用次数
+        call_count = 0
+        async def mock_execute(cmd):
+            nonlocal call_count
+            call_count += 1
+            return 5 # 模拟 get_security_count 返回值
+
+        mock_conn.execute.side_effect = mock_execute
+
+        # 设置非常短的心跳间隔以便测试
+        client = AsyncTdxClient("127.0.0.1", heartbeat_interval=0.1)
+        await client.connect()
+        
+        # 等待几次心跳周期
+        await asyncio.sleep(0.35)
+        
+        await client.close()
+        
+        # 0.35s 应该触发约 3 次心跳 (0.1, 0.2, 0.3)
+        assert call_count >= 3
+
+
+@pytest.mark.asyncio
+async def test_heartbeat_stops_on_close():
+    with patch("xmtdx.client.AsyncTdxConnection") as mock_conn_cls:
+        mock_conn = mock_conn_cls.return_value
+        mock_conn.connect = AsyncMock()
+        mock_conn.close = AsyncMock()
+        mock_conn.execute = AsyncMock(return_value=5)
+        
+        client = AsyncTdxClient("127.0.0.1", heartbeat_interval=0.01)
+        await client.connect()
+        assert client._heartbeat_task is not None
+        
+        task = client._heartbeat_task
+        await client.close()
+        
+        assert client._heartbeat_task is None
+        assert task.done() or task.cancelled()
+
+
+if __name__ == "__main__":
+    # 手动跑一下
+    async def run():
+        await test_heartbeat_sends_periodically()
+        await test_heartbeat_stops_on_close()
+        print("Heartbeat tests passed!")
+    
+    asyncio.run(run())