Add test for claude tools to_dict

getsentry · Jul 17, 2024 · a3bad9e · a3bad9e
1 parent b37aeeb
commit a3bad9e
Showing 1 changed file with 42 additions and 17 deletions.
diff --git a/tests/automation/agent/test_agent_tools.py b/tests/automation/agent/test_agent_tools.py
@@ -66,22 +66,47 @@ def test_chained_exception_handling(self, function_tool):
  assert result == expected
  mock_logger.exception.assert_called_once()
 
- def test_to_dict(self, function_tool):
- expected = {
- "type": "function",
- "function": {
- "name": "test_tool",
- "description": "A test tool",
- "parameters": {
- "type": "object",
- "properties": {
- "param1": {
- "type": "string",
- "description": "",
- }
+ @pytest.mark.parametrize(
+ "model, expected",
+ [
+ (
+ "gpt",
+ {
+ "type": "function",
+ "function": {
+ "name": "test_tool",
+ "description": "A test tool",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "param1": {
+ "type": "string",
+ "description": "",
+ }
+ },
+ "required": [],
+ },
  },
- "required": [],
  },
- },
- }
- assert function_tool.to_dict() == expected
+ ),
+ (
+ "claude",
+ {
+ "name": "test_tool",
+ "description": "A test tool",
+ "input_schema": {
+ "type": "object",
+ "properties": {
+ "param1": {
+ "type": "string",
+ "description": "",
+ }
+ },
+ "required": [],
+ },
+ },
+ ),
+ ],
+ )
+ def test_to_dict(self, function_tool, model, expected):
+ assert function_tool.to_dict(model=model) == expected