Coverage for tests/test_registry.py: 91%

1"""

2工具注册模块测试

3

4测试覆盖：

5- 工具注册装饰器

6- Schema 生成

7- 参数校验

8- 工具执行

9"""

10 import pytest

11 import json

12 from pydantic import BaseModel, Field

13 from qrclaw.tools.registry import (

14 register,

15 execute,

16 get_schemas,

17 need_confirm,

18 _tools,

19)

20

21

22# 测试用的参数模型

23 class TestArgs(BaseModel):

24 text: str = Field(description="测试文本")

25 count: int = Field(default=1, description="重复次数")

26

27

28 class TestTool:

29 """测试工具管理"""

30

31 def setup_method(self):

32 """每个测试前清空工具注册表"""

33 _tools.clear()

34

35 def test_register_tool(self):

36 """测试注册工具"""

37 @register(description="测试工具", args_model=TestArgs)

38 def test_tool(text: str, count: int = 1) -> str:

39 return text * count

40

41 # 验证注册成功

42 assert "test_tool" in _tools

43 assert _tools["test_tool"]["fn"] == test_tool

44 assert _tools["test_tool"]["confirm"] is False

45

46 def test_register_tool_with_confirm(self):

47 """测试注册需要确认的工具"""

48 @register(description="高风险工具", args_model=TestArgs, confirm=True)

49 def dangerous_tool(text: str, count: int = 1) -> str:

50 return text * count

51

52 assert need_confirm("dangerous_tool") is True

53

54 def test_need_confirm_default(self):

55 """测试默认不需要确认"""

56 @register(description="测试工具", args_model=TestArgs)

57 def normal_tool(text: str, count: int = 1) -> str:

58 return text * count

59

60 assert need_confirm("normal_tool") is False

61

62 def test_need_confirm_nonexistent(self):

63 """测试不存在的工具不需要确认"""

64 assert need_confirm("nonexistent_tool") is False

65

66

67 class TestSchema:

68 """测试 Schema 生成"""

69

70 def setup_method(self):

71 """每个测试前清空工具注册表"""

72 _tools.clear()

73

74 def test_get_schemas(self):

75 """测试获取所有工具 schema"""

76 @register(description="工具1", args_model=TestArgs)

77 def tool1(text: str, count: int = 1) -> str:

78 return text

79

80 @register(description="工具2", args_model=TestArgs)

81 def tool2(text: str, count: int = 1) -> str:

82 return text

83

84 schemas = get_schemas()

85

86 assert len(schemas) == 2

87 assert any(s["function"]["name"] == "tool1" for s in schemas)

88 assert any(s["function"]["name"] == "tool2" for s in schemas)

89

90 def test_schema_structure(self):

91 """测试 schema 结构"""

92 @register(description="测试工具", args_model=TestArgs)

93 def test_tool(text: str, count: int = 1) -> str:

94 return text

95

96 schema = _tools["test_tool"]["schema"]

97

98 assert schema["type"] == "function"

99 assert schema["function"]["name"] == "test_tool"

100 assert schema["function"]["description"] == "测试工具"

101 assert "parameters" in schema["function"]

102 assert "properties" in schema["function"]["parameters"]

103 assert "text" in schema["function"]["parameters"]["properties"]

104 assert "count" in schema["function"]["parameters"]["properties"]

105

106 def test_schema_required_fields(self):

107 """测试 required 字段"""

108 @register(description="测试工具", args_model=TestArgs)

109 def test_tool(text: str, count: int = 1) -> str:

110 return text

111

112 schema = _tools["test_tool"]["schema"]

113 params = schema["function"]["parameters"]

114

115 # text 是必填，count 有默认值

116 assert "text" in params.get("required", [])

117 assert "count" not in params.get("required", [])

118

119

120 class NoArgs(BaseModel):

121 """无参数工具的模型"""

122 pass

123

124

125 class TestExecution:

126 """测试工具执行"""

127

128 def setup_method(self):

129 """每个测试前清空工具注册表"""

130 _tools.clear()

131

132 def test_execute_tool_success(self):

133 """测试成功执行工具"""

134 @register(description="测试工具", args_model=TestArgs)

135 def test_tool(text: str, count: int = 1) -> str:

136 return text * count

137

138 result = execute("test_tool", json.dumps({"text": "hello", "count": 3}))

139

140 assert result == "hellohellohello"

141

142 def test_execute_tool_default_args(self):

143 """测试使用默认参数执行"""

144 @register(description="测试工具", args_model=TestArgs)

145 def test_tool(text: str, count: int = 1) -> str:

146 return text * count

147

148 result = execute("test_tool", json.dumps({"text": "hi"}))

149

150 assert result == "hi" # count 使用默认值 1

151

152 def test_execute_tool_invalid_json(self):

153 """测试无效 JSON 参数"""

154 @register(description="测试工具", args_model=TestArgs)

155 def test_tool(text: str, count: int = 1) -> str:

156 return text

157

158 result = execute("test_tool", "not a json")

159

160 assert "错误" in result

161 assert "JSON" in result

162

163 def test_execute_tool_missing_required(self):

164 """测试缺少必填参数"""

165 @register(description="测试工具", args_model=TestArgs)

166 def test_tool(text: str, count: int = 1) -> str:

167 return text

168

169 result = execute("test_tool", json.dumps({}))

170

171 assert "错误" in result or "校验失败" in result

172

173 def test_execute_tool_wrong_type(self):

174 """测试参数类型错误"""

175 @register(description="测试工具", args_model=TestArgs)

176 def test_tool(text: str, count: int = 1) -> str:

177 return text * count

178

179 result = execute("test_tool", json.dumps({"text": "hi", "count": "not a number"}))

180

181 assert "错误" in result or "校验失败" in result

182

183 def test_execute_nonexistent_tool(self):

184 """测试执行不存在的工具"""

185 result = execute("nonexistent", json.dumps({}))

186

187 assert "错误" in result

188 assert "找不到工具" in result

189

190 def test_execute_no_args_tool(self):

191 """测试无参数工具"""

192 @register(description="无参数工具", args_model=NoArgs)

193 def no_args_tool() -> str:

194 return "success"

195

196 result = execute("no_args_tool", json.dumps({}))

197

198 assert result == "success"

199

200 def test_execute_tool_exception(self):

201 """测试工具执行异常"""

202 @register(description="会报错的工具", args_model=TestArgs)

203 def error_tool(text: str, count: int = 1) -> str:

204 raise ValueError("故意报错")

205

206 result = execute("error_tool", json.dumps({"text": "test"}))

207

208 # 应该捕获异常并返回错误信息

209 assert "错误" in result or "校验失败" in result