5bc275b81c
\`get_chat_images\` 和 \`get_voice_messages\` 仅有 \`limit\`, 接口与 \`get_chat_history\` / \`search_messages\` (\`offset\` + \`start_time\` + \`end_time\`) 不对齐: 1. 查不了"某段时间内的图片/语音" 2. 不支持分页, 单次取 \`limit=1000\` 一次性拉 3. LLM 用同样模式调不同工具时签名不一致, 容易出错 两个工具各加 3 个可选参数: - \`offset: int = 0\` - \`start_time: str = ""\` - \`end_time: str = ""\` 复用上游已有的 \`_validate_pagination\` + \`_parse_time_range\` helpers。 - 输入校验失败立即报错 (offset 负数 / 时间格式错 / start > end) - 每 shard 拉 \`limit + offset\` 张候选, 合并后全局 \`create_time DESC\` 排序, 切片 \`[offset : offset + limit]\` 出本页 - 单 shard 凑得起本页, 避免某 shard 缺数据时本页变短 - header 显示 offset/limit 和时间范围 (传了才显示) - 加 \`start_ts=None\` / \`end_ts=None\` 参数 - SQL 动态拼 \`create_time >= ?\` / \`<= ?\` clause - 不传时序参数完全等价旧行为 (向后兼容) - 同样 3 个参数 + \`_validate_pagination\` + \`_parse_time_range\` - VoiceInfo 表 SQL 动态拼 \`chat_name_id = ? AND create_time ?...\` - 多 shard 各取 \`limit + offset\` 后合并切片 新增 \`tests/test_chat_images_query_align.py\` 8 个 case: - offset 负数报错 - start > end 报错 - candidate_limit = limit + offset (shard 调用确认) - 时间参数正确解析为 unix 秒并透传 - offset=2, limit=2 切到全局排序后第 3-4 张 - header 包含时间范围 - header 包含 offset/limit - 默认调用(不传新参)行为与旧接口一致 修改 \`tests/test_get_chat_images_multishard.py\` 的 fake_list 签名: - 旧: \`(db_path, table_name, username, lim)\` 位置参 - 新: \`(db_path, table_name, username, limit=20, start_ts=None, end_ts=None)\` - 既支持旧调用模式 (kwargs), 也兼容新签名 全量 \`pytest tests/\` 208/208 通过。 3 个可选参数全部带默认值 → 既有调用方零修改。 shard candidate=\`limit+offset\` 的成本: 大 offset 时单 shard 请求量 增大。但 image/voice 表每 chat 单 shard 一般 < 10K 条, 实际 cost 可 忽略。如果将来要做"翻 100 页"级深翻, 可以加 keyset pagination, 现 在 offset 模式与 \`get_chat_history\` 一致即可。 与 #103 / #104 触碰同一文件, 合并顺序无所谓 — 后合的 rebase 即可。
98 lines
3.5 KiB
Python
98 lines
3.5 KiB
Python
"""测试 get_chat_images 新增的 offset / start_time / end_time 参数。"""
|
|
import os
|
|
import sys
|
|
from unittest.mock import patch
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
import mcp_server
|
|
|
|
|
|
def _img(local_id, create_time, md5=None, size=None):
|
|
info = {'local_id': local_id, 'create_time': create_time, 'md5': md5}
|
|
if size is not None:
|
|
info['size'] = size
|
|
return info
|
|
|
|
|
|
def _run_with(shard_images_map, **kwargs):
|
|
"""Helper: stub collaborators and call get_chat_images with new kwargs."""
|
|
shards = [{'db_path': k, 'table_name': 'Msg_x'} for k in shard_images_map]
|
|
captured = {'calls': []}
|
|
|
|
def fake_list(db_path, table_name, username, limit=20, start_ts=None, end_ts=None):
|
|
captured['calls'].append({
|
|
'db_path': db_path, 'limit': limit, 'start_ts': start_ts, 'end_ts': end_ts,
|
|
})
|
|
return shard_images_map.get(db_path, [])
|
|
|
|
with patch.object(mcp_server, 'resolve_username', return_value='wxid_demo'), \
|
|
patch.object(mcp_server, 'get_contact_names', return_value={'wxid_demo': 'Demo'}), \
|
|
patch.object(mcp_server, '_find_msg_tables_for_user', return_value=shards), \
|
|
patch.object(mcp_server._image_resolver, 'list_chat_images', side_effect=fake_list):
|
|
return mcp_server.get_chat_images('Demo', **kwargs), captured
|
|
|
|
|
|
def test_invalid_offset_returns_error():
|
|
out, _ = _run_with({}, offset=-1)
|
|
assert '错误' in out
|
|
|
|
|
|
def test_invalid_time_range_returns_error():
|
|
"""start_time 晚于 end_time 应报错。"""
|
|
out, _ = _run_with({}, start_time='2026-05-10', end_time='2026-05-01')
|
|
assert '错误' in out
|
|
|
|
|
|
def test_candidate_limit_includes_offset():
|
|
"""每 shard 拉 limit+offset 张候选, 保证全局分页能切到正确的页。"""
|
|
_, captured = _run_with({'/a': [_img(1, 1000)]}, limit=5, offset=10)
|
|
assert captured['calls'][0]['limit'] == 15
|
|
|
|
|
|
def test_start_end_ts_forwarded_to_shard_query():
|
|
"""start_time / end_time 解析为 unix 秒后透传给 shard 查询。"""
|
|
_, captured = _run_with(
|
|
{'/a': []},
|
|
start_time='2026-05-01',
|
|
end_time='2026-05-31',
|
|
)
|
|
call = captured['calls'][0]
|
|
assert call['start_ts'] is not None
|
|
assert call['end_ts'] is not None
|
|
assert call['start_ts'] < call['end_ts']
|
|
|
|
|
|
def test_offset_slices_paged_window():
|
|
"""offset=2, limit=2 取全局排序后第 3-4 张图片。"""
|
|
shard_a = [_img(1, 1100), _img(2, 1000)]
|
|
shard_b = [_img(3, 1300), _img(4, 1200)]
|
|
out, _ = _run_with({'/a': shard_a, '/b': shard_b}, limit=2, offset=2)
|
|
# 全局排序后顺序: 1300, 1200, 1100, 1000 → 第 3-4 是 1100, 1000 → local_id 1, 2
|
|
assert 'local_id=1' in out
|
|
assert 'local_id=2' in out
|
|
assert 'local_id=3' not in out
|
|
assert 'local_id=4' not in out
|
|
|
|
|
|
def test_header_shows_time_range_when_given():
|
|
shard_a = [_img(1, 1000, md5='abc')]
|
|
out, _ = _run_with({'/a': shard_a}, start_time='2026-05-01')
|
|
assert '时间范围' in out
|
|
assert '2026-05-01' in out
|
|
|
|
|
|
def test_header_shows_offset_limit():
|
|
shard_a = [_img(1, 1000, md5='abc')]
|
|
out, _ = _run_with({'/a': shard_a}, limit=10, offset=20)
|
|
assert 'offset=20' in out
|
|
assert 'limit=10' in out
|
|
|
|
|
|
def test_default_behavior_unchanged():
|
|
"""不传新参数时行为与旧接口一致 — offset=0 切片就是 [:limit]。"""
|
|
shard_a = [_img(1, 1100, md5='a1'), _img(2, 1000, md5='a2')]
|
|
out, _ = _run_with({'/a': shard_a})
|
|
assert 'local_id=1' in out
|
|
assert 'local_id=2' in out
|