Coverage for agentos/tests/test_1_1_4_features.py: 79%

1"""v1.1.4 新特性集成测试。""" 

2 

3from __future__ import annotations 

4 

5import pytest 

6 

7# ══════════════════════════════════════════════════════════════════════════════ 

8# ToolRiskRating 测试 

9# ══════════════════════════════════════════════════════════════════════════════ 

10 

11 

12class TestToolRiskRating: 

13 def test_risk_level_enum(self): 

14 from agentos.tools.risk import ToolRiskLevel 

15 

16 assert ToolRiskLevel.LOW.value == "low" 

17 assert ToolRiskLevel.CRITICAL.value == "critical" 

18 assert len(list(ToolRiskLevel)) == 4 

19 

20 def test_risk_rating_defaults(self): 

21 from agentos.tools.risk import ToolRiskRating 

22 

23 r = ToolRiskRating() 

24 assert r.level.value == "medium" 

25 assert r.reversible is True 

26 assert r.requires_approval is False 

27 assert r.requires_user_confirm() is False 

28 

29 def test_requires_confirm_high(self): 

30 from agentos.tools.risk import ToolRiskLevel, ToolRiskRating 

31 

32 r = ToolRiskRating(level=ToolRiskLevel.HIGH) 

33 assert r.requires_user_confirm() is True 

34 

35 def test_requires_confirm_critical(self): 

36 from agentos.tools.risk import ToolRiskLevel, ToolRiskRating 

37 

38 r = ToolRiskRating(level=ToolRiskLevel.CRITICAL) 

39 assert r.requires_user_confirm() is True 

40 

41 def test_requires_confirm_financial(self): 

42 from agentos.tools.risk import ToolRiskRating 

43 

44 r = ToolRiskRating(financial_impact=True) 

45 assert r.requires_user_confirm() is True 

46 

47 def test_get_risk_preset_list_files(self): 

48 from agentos.tools.risk import ToolRiskLevel, get_risk_preset 

49 

50 r = get_risk_preset("list_files") 

51 assert r is not None 

52 assert r.level == ToolRiskLevel.LOW 

53 

54 def test_get_risk_preset_delete_file(self): 

55 from agentos.tools.risk import ToolRiskLevel, get_risk_preset 

56 

57 r = get_risk_preset("delete_file") 

58 assert r is not None 

59 assert r.level == ToolRiskLevel.HIGH 

60 assert r.requires_approval is True 

61 

62 def test_get_risk_preset_payment(self): 

63 from agentos.tools.risk import ToolRiskLevel, get_risk_preset 

64 

65 r = get_risk_preset("execute_payment") 

66 assert r is not None 

67 assert r.level == ToolRiskLevel.CRITICAL 

68 assert r.financial_impact is True 

69 

70 def test_get_risk_preset_case_insensitive(self): 

71 from agentos.tools.risk import get_risk_preset 

72 

73 assert get_risk_preset("DELETE_FILE") is not None 

74 

75 def test_infer_risk_level_keyword_delete(self): 

76 from agentos.tools.risk import ToolRiskLevel, infer_risk_level 

77 

78 r = infer_risk_level("purge_records", "delete all records") 

79 assert r.level == ToolRiskLevel.HIGH 

80 

81 def test_infer_risk_level_keyword_write(self): 

82 from agentos.tools.risk import ToolRiskLevel, infer_risk_level 

83 

84 r = infer_risk_level("update_profile") 

85 assert r.level == ToolRiskLevel.MEDIUM 

86 

87 def test_infer_risk_level_default(self): 

88 from agentos.tools.risk import ToolRiskLevel, infer_risk_level 

89 

90 r = infer_risk_level("get_status") 

91 assert r.level == ToolRiskLevel.LOW 

92 

93 

94# ══════════════════════════════════════════════════════════════════════════════ 

95# Middleware Pipeline 测试 

96# ══════════════════════════════════════════════════════════════════════════════ 

97 

98 

99class TestMiddlewarePipeline: 

100 @pytest.mark.asyncio 

101 async def test_empty_pipeline_allows(self): 

102 from agentos.core.middleware import MiddlewareContext, MiddlewarePhase, MiddlewarePipeline 

103 

104 pipe = MiddlewarePipeline() 

105 ctx = MiddlewareContext(phase=MiddlewarePhase.PRE_LLM, prompt="hello") 

106 decision = await pipe.pre_llm(ctx) 

107 assert decision.allow is True 

108 

109 @pytest.mark.asyncio 

110 async def test_blocking_middleware(self): 

111 from agentos.core.middleware import ( 

112 AgentMiddleware, 

113 MiddlewareContext, 

114 MiddlewareDecision, 

115 MiddlewarePhase, 

116 MiddlewarePipeline, 

117 ) 

118 

119 class Blocker(AgentMiddleware): 

120 name = "blocker" 

121 

122 @property 

123 def phases(self): 

124 return [MiddlewarePhase.PRE_LLM] 

125 

126 async def process(self, ctx): 

127 return MiddlewareDecision(allow=False, reason="blocked by test", action="block") 

128 

129 pipe = MiddlewarePipeline([Blocker()]) 

130 ctx = MiddlewareContext(phase=MiddlewarePhase.PRE_LLM, prompt="test") 

131 decision = await pipe.pre_llm(ctx) 

132 assert decision.allow is False 

133 assert "blocked by test" in decision.reason 

134 

135 @pytest.mark.asyncio 

136 async def test_transform_middleware(self): 

137 from agentos.core.middleware import ( 

138 AgentMiddleware, 

139 MiddlewareContext, 

140 MiddlewareDecision, 

141 MiddlewarePhase, 

142 MiddlewarePipeline, 

143 ) 

144 

145 class UpperCaseTransform(AgentMiddleware): 

146 name = "upper" 

147 

148 @property 

149 def phases(self): 

150 return [MiddlewarePhase.PRE_LLM] 

151 

152 async def process(self, ctx): 

153 if ctx.prompt: 

154 new_ctx = MiddlewareContext(**{**ctx.__dict__}) 

155 new_ctx.prompt = ctx.prompt.upper() 

156 return MiddlewareDecision( 

157 allow=True, action="transform", modified_context=new_ctx 

158 ) 

159 return MiddlewareDecision(allow=True) 

160 

161 pipe = MiddlewarePipeline([UpperCaseTransform()]) 

162 ctx = MiddlewareContext(phase=MiddlewarePhase.PRE_LLM, prompt="hello") 

163 decision = await pipe.pre_llm(ctx) 

164 assert decision.allow is True 

165 assert decision.modified_context is not None 

166 assert decision.modified_context.prompt == "HELLO" 

167 

168 @pytest.mark.asyncio 

169 async def test_chain_add(self): 

170 from agentos.core.middleware import MiddlewarePipeline 

171 

172 pipe = MiddlewarePipeline() 

173 from agentos.core.middleware import AuditLogMiddleware 

174 

175 pipe.add(AuditLogMiddleware()) 

176 assert "audit_log" in pipe.middleware_names 

177 

178 @pytest.mark.asyncio 

179 async def test_remove(self): 

180 from agentos.core.middleware import AuditLogMiddleware, MiddlewarePipeline 

181 

182 pipe = MiddlewarePipeline([AuditLogMiddleware()]) 

183 pipe.remove("audit_log") 

184 assert "audit_log" not in pipe.middleware_names 

185 

186 @pytest.mark.asyncio 

187 async def test_phase_filtering(self): 

188 from agentos.core.middleware import ( 

189 MiddlewareContext, 

190 MiddlewarePhase, 

191 MiddlewarePipeline, 

192 PIIMaskingMiddleware, 

193 ) 

194 

195 pipe = MiddlewarePipeline([PIIMaskingMiddleware()]) 

196 # PIIMaskingMiddleware only listens on PRE_LLM 

197 ctx = MiddlewareContext(phase=MiddlewarePhase.PRE_TOOL, tool_name="test") 

198 decision = await pipe.pre_tool(ctx) 

199 assert decision.allow is True # It should pass through since no middleware listens 

200 

201 

202# ══════════════════════════════════════════════════════════════════════════════ 

203# Enhanced CostTracker + RunCostSession 测试 

204# ══════════════════════════════════════════════════════════════════════════════ 

205 

206 

207class TestRunCostSession: 

208 def test_session_lifecycle(self): 

209 import time 

210 

211 from agentos.cost.tracker import RunCostSession 

212 

213 session = RunCostSession(run_id="test-123") 

214 assert session.run_id == "test-123" 

215 assert session.call_count == 0 

216 assert session.total_cost == 0 

217 

218 # Record some usage 

219 from agentos.cost.tracker import UsageRecord 

220 

221 session.records.append( 

222 UsageRecord( 

223 model="deepseek-v3.1", 

224 input_tokens=1000, 

225 output_tokens=500, 

226 cost_usd=0.01, 

227 run_id="test-123", 

228 ) 

229 ) 

230 session.records.append( 

231 UsageRecord( 

232 model="deepseek-v3.1", 

233 input_tokens=2000, 

234 output_tokens=800, 

235 cost_usd=0.02, 

236 run_id="test-123", 

237 ) 

238 ) 

239 session.finished_at = time.time() 

240 

241 assert session.call_count == 2 

242 assert session.total_cost == 0.03 

243 assert session.duration_seconds > 0 

244 

245 def test_total_tokens(self): 

246 from agentos.cost.tracker import RunCostSession, UsageRecord 

247 

248 session = RunCostSession(run_id="t") 

249 session.records.append( 

250 UsageRecord(model="m", input_tokens=100, output_tokens=50, cost_usd=0.0) 

251 ) 

252 session.records.append( 

253 UsageRecord(model="m", input_tokens=200, output_tokens=100, cost_usd=0.0) 

254 ) 

255 assert session.total_tokens == {"input": 300, "output": 150, "total": 450} 

256 

257 

258class TestCostTrackerEnhanced: 

259 def test_start_end_session(self): 

260 from agentos.cost.tracker import CostTracker 

261 

262 tracker = CostTracker() 

263 rid = tracker.start_session() 

264 assert len(tracker.active_sessions) == 1 

265 session = tracker.end_session(rid) 

266 assert session is not None 

267 assert session.finished_at is not None 

268 assert len(tracker.active_sessions) == 0 

269 

270 def test_record_with_session(self): 

271 from agentos.cost.tracker import CostTracker 

272 

273 tracker = CostTracker() 

274 rid = tracker.start_session() 

275 tracker.record( 

276 "deepseek-v3.1", {"prompt_tokens": 1000, "completion_tokens": 500}, run_id=rid 

277 ) 

278 tracker.record( 

279 "deepseek-v3.1", {"prompt_tokens": 500, "completion_tokens": 200}, run_id=rid 

280 ) 

281 session = tracker.end_session(rid) 

282 assert session.call_count == 2 

283 assert session.total_cost > 0 

284 

285 def test_cost_by_session(self): 

286 from agentos.cost.tracker import CostTracker 

287 

288 tracker = CostTracker() 

289 r1 = tracker.start_session() 

290 r2 = tracker.start_session() 

291 tracker.record( 

292 "deepseek-v3.1", {"prompt_tokens": 1000, "completion_tokens": 100}, run_id=r1 

293 ) 

294 tracker.record("deepseek-v3.1", {"prompt_tokens": 500, "completion_tokens": 50}, run_id=r2) 

295 costs = tracker.cost_by_session() 

296 assert r1 in costs 

297 assert r2 in costs 

298 assert costs[r2] < costs[r1] 

299 

300 def test_get_session_active_and_completed(self): 

301 from agentos.cost.tracker import CostTracker 

302 

303 tracker = CostTracker() 

304 rid = tracker.start_session() 

305 assert tracker.get_session(rid) is not None 

306 tracker.end_session(rid) 

307 assert tracker.get_session(rid) is not None # Should find completed 

308 

309 def test_record_with_cache(self): 

310 from agentos.cost.tracker import CostTracker 

311 

312 tracker = CostTracker() 

313 cost = tracker.record_with_cache("deepseek-v3.1", 1000, 500) 

314 assert cost > 0 

315 assert tracker.total_cost == cost