Coverage for tests / test_session.py: 100%

1""" 

2Session 会话管理模块测试 

3 

4测试覆盖： 

5- 会话创建与恢复 

6- 消息添加与持久化 

7- Token 计算 

8- 计划管理 

9""" 

10import pytest 

11import json 

12from pathlib import Path 

13from qrclaw.memory.session import ( 

14 Session, 

15 list_sessions, 

16 get_last_session_id, 

17 delete_session, 

18 count_tokens, 

19) 

20 

21 

22class TestSessionCreation: 

23 """测试会话创建""" 

24 

25 def test_create_new_session(self, temp_dir): 

26 """测试创建新会话""" 

27 session = Session(temp_dir, resume=False) 

28 

29 # 验证会话 ID 格式 

30 assert session.session_id is not None 

31 assert len(session.session_id) > 0 

32 

33 # 验证初始状态 

34 assert session.messages == [] 

35 assert session.prompt_tokens == 0 

36 assert session.active_plan is None 

37 

38 def test_create_session_with_id(self, temp_dir): 

39 """测试指定 ID 创建会话""" 

40 session = Session(temp_dir, session_id="test-session-123") 

41 

42 assert session.session_id == "test-session-123" 

43 assert session._path.name == "test-session-123.json" 

44 

45 def test_resume_last_session(self, temp_dir): 

46 """测试恢复最近会话""" 

47 # 创建第一个会话 

48 session1 = Session(temp_dir, resume=False) 

49 session1.add({"role": "user", "content": "hello"}) 

50 

51 # 创建第二个会话（应该恢复第一个） 

52 session2 = Session(temp_dir, resume=True) 

53 

54 # 应该是同一个会话 

55 assert session2.session_id == session1.session_id 

56 assert len(session2.messages) == 1 

57 assert session2.messages[0]["content"] == "hello" 

58 

59 def test_no_resume_creates_new(self, temp_dir): 

60 """测试不恢复时创建新会话""" 

61 # 创建第一个会话 

62 session1 = Session(temp_dir, resume=False) 

63 session1.add({"role": "user", "content": "test"}) 

64 

65 # 不恢复，创建新会话 

66 session2 = Session(temp_dir, resume=False) 

67 

68 # 应该是不同的会话 

69 assert session2.session_id != session1.session_id 

70 assert len(session2.messages) == 0 

71 

72 

73class TestSessionMessages: 

74 """测试消息管理""" 

75 

76 def test_add_message(self, temp_dir): 

77 """测试添加消息""" 

78 session = Session(temp_dir, resume=False) 

79 

80 session.add({"role": "user", "content": "你好"}) 

81 session.add({"role": "assistant", "content": "你好！有什么可以帮你的？"}) 

82 

83 assert len(session.messages) == 2 

84 assert session.messages[0]["role"] == "user" 

85 assert session.messages[0]["content"] == "你好" 

86 

87 # 验证持久化 

88 saved_data = json.loads(session._path.read_text(encoding="utf-8")) 

89 assert len(saved_data) == 2 

90 

91 def test_message_persistence(self, temp_dir): 

92 """测试消息持久化""" 

93 # 创建会话并添加消息 

94 session_id = "persist-test" 

95 session1 = Session(temp_dir, session_id=session_id) 

96 session1.add({"role": "user", "content": "测试持久化"}) 

97 session1.add({"role": "assistant", "content": "收到"}) 

98 

99 # 重新加载会话 

100 session2 = Session(temp_dir, session_id=session_id) 

101 

102 # 验证消息已加载 

103 assert len(session2.messages) == 2 

104 assert session2.messages[0]["content"] == "测试持久化" 

105 

106 def test_clear_session(self, temp_dir): 

107 """测试清空会话""" 

108 session = Session(temp_dir, resume=False) 

109 session.add({"role": "user", "content": "test"}) 

110 

111 # 清空 

112 session.clear() 

113 

114 assert session.messages == [] 

115 assert not session._path.exists() 

116 

117 

118class TestTokenCounting: 

119 """测试 Token 计算""" 

120 

121 def test_count_tokens_basic(self): 

122 """测试基础 token 计算""" 

123 messages = [ 

124 {"role": "user", "content": "hello"} 

125 ] 

126 tokens = count_tokens(messages) 

127 

128 # 至少包含消息内容 token 

129 assert tokens > 0 

130 

131 def test_count_tokens_multiple_messages(self): 

132 """测试多条消息的 token 计算""" 

133 messages = [ 

134 {"role": "user", "content": "你好"}, 

135 {"role": "assistant", "content": "你好！有什么可以帮你的？"} 

136 ] 

137 tokens = count_tokens(messages) 

138 

139 # 应该比单条消息多 

140 single_tokens = count_tokens([messages[0]]) 

141 assert tokens > single_tokens 

142 

143 def test_count_tokens_empty_messages(self): 

144 """测试空消息列表""" 

145 tokens = count_tokens([]) 

146 # 空消息应该只有基础开销 

147 assert tokens == 2 # 对话开销 

148 

149 

150class TestSessionPlan: 

151 """测试计划管理""" 

152 

153 def test_set_plan(self, temp_dir): 

154 """测试设置计划""" 

155 session = Session(temp_dir, resume=False) 

156 

157 steps = [ 

158 {"id": 1, "description": "步骤1"}, 

159 {"id": 2, "description": "步骤2"} 

160 ] 

161 session.set_plan("测试目标", steps) 

162 

163 assert session.active_plan is not None 

164 assert session.active_plan["goal"] == "测试目标" 

165 assert len(session.active_plan["steps"]) == 2 

166 assert session.active_plan["steps"][0]["done"] is False 

167 

168 def test_complete_step(self, temp_dir): 

169 """测试完成步骤""" 

170 session = Session(temp_dir, resume=False) 

171 session.set_plan("目标", [{"id": 1, "description": "步骤1"}]) 

172 

173 # 完成步骤前验证状态 

174 assert session.active_plan["steps"][0]["done"] is False 

175 

176 # 完成步骤 

177 all_done = session.complete_step(1) 

178 

179 # 所有步骤完成后计划会被清空 

180 assert all_done is True 

181 assert session.active_plan is None 

182 

183 def test_complete_step_partial(self, temp_dir): 

184 """测试部分完成步骤""" 

185 session = Session(temp_dir, resume=False) 

186 steps = [ 

187 {"id": 1, "description": "步骤1"}, 

188 {"id": 2, "description": "步骤2"} 

189 ] 

190 session.set_plan("目标", steps) 

191 

192 # 完成第一个步骤 

193 all_done = session.complete_step(1) 

194 

195 assert session.active_plan["steps"][0]["done"] is True 

196 assert session.active_plan["steps"][1]["done"] is False 

197 assert all_done is False # 还有未完成的步骤 

198 

199 def test_clear_plan(self, temp_dir): 

200 """测试清空计划""" 

201 session = Session(temp_dir, resume=False) 

202 session.set_plan("目标", [{"id": 1, "description": "步骤1"}]) 

203 

204 session.clear_plan() 

205 

206 assert session.active_plan is None 

207 

208 

209class TestSessionList: 

210 """测试会话列表管理""" 

211 

212 def test_list_sessions(self, temp_dir): 

213 """测试列出会话""" 

214 # 创建多个会话 

215 Session(temp_dir, session_id="session-1").add({"role": "user", "content": "test1"}) 

216 Session(temp_dir, session_id="session-2").add({"role": "user", "content": "test2"}) 

217 

218 sessions = list_sessions(temp_dir) 

219 

220 assert len(sessions) == 2 

221 assert any(s["id"] == "session-1" for s in sessions) 

222 assert any(s["id"] == "session-2" for s in sessions) 

223 

224 def test_get_last_session_id(self, temp_dir): 

225 """测试获取最近会话 ID""" 

226 # 创建两个会话 

227 Session(temp_dir, session_id="older").add({"role": "user", "content": "test1"}) 

228 

229 # 等一下，确保时间戳不同 

230 import time 

231 time.sleep(0.01) 

232 

233 Session(temp_dir, session_id="newer").add({"role": "user", "content": "test2"}) 

234 

235 # 应该返回最新的 

236 last_id = get_last_session_id(temp_dir) 

237 assert last_id == "newer" 

238 

239 def test_delete_session(self, temp_dir): 

240 """测试删除会话""" 

241 session = Session(temp_dir, session_id="to-delete") 

242 session.add({"role": "user", "content": "test"}) 

243 

244 # 删除 

245 result = delete_session("to-delete", temp_dir) 

246 

247 assert result is True 

248 assert not session._path.exists() 

249 

250 def test_delete_nonexistent_session(self, temp_dir): 

251 """测试删除不存在的会话""" 

252 result = delete_session("nonexistent", temp_dir) 

253 assert result is False