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

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 ( 

103 MiddlewareContext, 

104 MiddlewarePhase, 

105 MiddlewarePipeline, 

106 ) 

107 

108 pipe = MiddlewarePipeline() 

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

110 decision = await pipe.pre_llm(ctx) 

111 assert decision.allow is True 

112 

113 @pytest.mark.asyncio 

114 async def test_blocking_middleware(self): 

115 from agentos.core.middleware import ( 

116 AgentMiddleware, 

117 MiddlewareContext, 

118 MiddlewareDecision, 

119 MiddlewarePhase, 

120 MiddlewarePipeline, 

121 ) 

122 

123 class Blocker(AgentMiddleware): 

124 name = "blocker" 

125 

126 @property 

127 def phases(self): 

128 return [MiddlewarePhase.PRE_LLM] 

129 

130 async def process(self, ctx): 

131 return MiddlewareDecision( 

132 allow=False, reason="blocked by test", action="block" 

133 ) 

134 

135 pipe = MiddlewarePipeline([Blocker()]) 

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

137 decision = await pipe.pre_llm(ctx) 

138 assert decision.allow is False 

139 assert "blocked by test" in decision.reason 

140 

141 @pytest.mark.asyncio 

142 async def test_transform_middleware(self): 

143 from agentos.core.middleware import ( 

144 AgentMiddleware, 

145 MiddlewareContext, 

146 MiddlewareDecision, 

147 MiddlewarePhase, 

148 MiddlewarePipeline, 

149 ) 

150 

151 class UpperCaseTransform(AgentMiddleware): 

152 name = "upper" 

153 

154 @property 

155 def phases(self): 

156 return [MiddlewarePhase.PRE_LLM] 

157 

158 async def process(self, ctx): 

159 if ctx.prompt: 

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

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

162 return MiddlewareDecision( 

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

164 ) 

165 return MiddlewareDecision(allow=True) 

166 

167 pipe = MiddlewarePipeline([UpperCaseTransform()]) 

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

169 decision = await pipe.pre_llm(ctx) 

170 assert decision.allow is True 

171 assert decision.modified_context is not None 

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

173 

174 @pytest.mark.asyncio 

175 async def test_chain_add(self): 

176 from agentos.core.middleware import MiddlewarePipeline 

177 

178 pipe = MiddlewarePipeline() 

179 from agentos.core.middleware import AuditLogMiddleware 

180 

181 pipe.add(AuditLogMiddleware()) 

182 assert "audit_log" in pipe.middleware_names 

183 

184 @pytest.mark.asyncio 

185 async def test_remove(self): 

186 from agentos.core.middleware import AuditLogMiddleware, MiddlewarePipeline 

187 

188 pipe = MiddlewarePipeline([AuditLogMiddleware()]) 

189 pipe.remove("audit_log") 

190 assert "audit_log" not in pipe.middleware_names 

191 

192 @pytest.mark.asyncio 

193 async def test_phase_filtering(self): 

194 from agentos.core.middleware import ( 

195 MiddlewareContext, 

196 MiddlewarePhase, 

197 MiddlewarePipeline, 

198 PIIMaskingMiddleware, 

199 ) 

200 

201 pipe = MiddlewarePipeline([PIIMaskingMiddleware()]) 

202 # PIIMaskingMiddleware only listens on PRE_LLM 

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

204 decision = await pipe.pre_tool(ctx) 

205 assert ( 

206 decision.allow is True 

207 ) # It should pass through since no middleware listens 

208 

209 

210# ══════════════════════════════════════════════════════════════════════════════ 

211# Enhanced CostTracker + RunCostSession 测试 

212# ══════════════════════════════════════════════════════════════════════════════ 

213 

214 

215class TestRunCostSession: 

216 """Test CostTracker core functionality (record, report, budget).""" 

217 

218 def test_session_lifecycle(self): 

219 from agentos.cost.tracker import CostTracker 

220 

221 tracker = CostTracker() 

222 summary = tracker.report_dict() 

223 assert summary["total_cost"] == 0 

224 assert summary["total_tokens"] == 0 

225 

226 tracker.record("deepseek-v3.1", input_tokens=1000, output_tokens=500) 

227 summary = tracker.report_dict() 

228 assert summary["total_cost"] > 0 

229 assert summary["total_tokens"] == 1500 

230 

231 tracker.record("deepseek-v3.1", input_tokens=2000, output_tokens=800) 

232 summary = tracker.report_dict() 

233 assert summary["total_cost"] > 0 

234 assert summary["total_tokens"] == 4300 

235 

236 def test_total_tokens(self): 

237 from agentos.cost.tracker import CostTracker 

238 

239 tracker = CostTracker() 

240 tracker.record("gpt-4o", input_tokens=100, output_tokens=50) 

241 tracker.record("gpt-4o", input_tokens=200, output_tokens=100) 

242 summary = tracker.report_dict() 

243 assert summary["total_tokens"] == 450 

244 

245 

246class TestCostTrackerEnhanced: 

247 def test_record_and_report(self): 

248 from agentos.cost.tracker import CostTracker 

249 

250 tracker = CostTracker() 

251 tracker.record("deepseek-v3.1", input_tokens=1000, output_tokens=500) 

252 report = tracker.report_dict() 

253 assert report["total_cost"] > 0 

254 assert report["total_tokens"] == 1500 

255 

256 def test_record_multiple_models(self): 

257 from agentos.cost.tracker import CostTracker 

258 

259 tracker = CostTracker() 

260 tracker.record("gpt-4o", input_tokens=500, output_tokens=200) 

261 tracker.record("claude-3-5-sonnet", input_tokens=1000, output_tokens=500) 

262 report = tracker.report_dict() 

263 assert len(report["by_model"]) >= 2 

264 

265 def test_get_price_fallback(self): 

266 from agentos.cost.tracker import CostTracker 

267 

268 tracker = CostTracker() 

269 price = tracker.get_price("gpt-4o") 

270 assert price is not None 

271 

272 def test_reset(self): 

273 from agentos.cost.tracker import CostTracker 

274 

275 tracker = CostTracker() 

276 tracker.record("deepseek-v3.1", input_tokens=1000, output_tokens=500) 

277 report = tracker.report_dict() 

278 assert report["total_cost"] > 0 

279 tracker.reset() 

280 report = tracker.report_dict() 

281 assert report["total_cost"] == 0 

282 assert report["total_tokens"] == 0 

283 

284 def test_budget_check(self): 

285 from agentos.cost.tracker import Budget, CostTracker 

286 

287 tracker = CostTracker(budgets=[Budget(name="test-budget", limit=0.01)]) 

288 tracker.record("deepseek-v3.1", input_tokens=10000, output_tokens=5000) 

289 alerts = tracker.check_budget() 

290 # May or may not exceed depending on pricing; verify method runs 

291 assert isinstance(alerts, list)