133 lines
2.9 KiB
Go
133 lines
2.9 KiB
Go
package main
|
|
|
|
import (
|
|
"grailbench/config"
|
|
"log/slog"
|
|
"os"
|
|
"testing"
|
|
)
|
|
|
|
// TestParserToolDetection tests the parser implementations with mock responses
|
|
func TestParserToolDetection(t *testing.T) {
|
|
// Create a logger for testing
|
|
logger := slog.New(slog.NewTextHandler(os.Stderr, nil))
|
|
|
|
// Test cases with different response types
|
|
testCases := []struct {
|
|
name string
|
|
responseJSON string
|
|
expectedOutput string
|
|
isToolCall bool
|
|
expectedTool string
|
|
}{
|
|
{
|
|
name: "RegularTextResponse",
|
|
responseJSON: `{
|
|
"id": "test-id",
|
|
"object": "text.completion",
|
|
"created": 1234567890,
|
|
"model": "deepseek-chat",
|
|
"choices": [
|
|
{
|
|
"text": "The capital of France is Paris.",
|
|
"index": 0,
|
|
"finish_reason": "stop"
|
|
}
|
|
]
|
|
}`,
|
|
expectedOutput: "The capital of France is Paris.",
|
|
isToolCall: false,
|
|
},
|
|
{
|
|
name: "ToolCallResponse",
|
|
responseJSON: `{
|
|
"id": "test-id",
|
|
"object": "chat.completion",
|
|
"created": 1234567890,
|
|
"model": "deepseek-chat",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"finish_reason": "tool_calls",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": null,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_123",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_current_timestamp",
|
|
"arguments": "{}"
|
|
}
|
|
}
|
|
]
|
|
}
|
|
}
|
|
]
|
|
}`,
|
|
expectedOutput: "[TOOL_CALL:get_current_timestamp]",
|
|
isToolCall: true,
|
|
expectedTool: "get_current_timestamp",
|
|
},
|
|
{
|
|
name: "RegularMessageResponse",
|
|
responseJSON: `{
|
|
"id": "test-id",
|
|
"object": "chat.completion",
|
|
"created": 1234567890,
|
|
"model": "deepseek-chat",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"finish_reason": "stop",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": "Hello, how can I help you today?"
|
|
}
|
|
}
|
|
]
|
|
}`,
|
|
expectedOutput: "Hello, how can I help you today?",
|
|
isToolCall: false,
|
|
},
|
|
}
|
|
|
|
// Test DeepSeekParser
|
|
t.Run("DeepSeekParser", func(t *testing.T) {
|
|
parser := NewDeepSeekParser(logger, config.LoadConfigOrDefault("config.toml"))
|
|
|
|
for _, tc := range testCases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
result, err := parser.ParseBytes([]byte(tc.responseJSON))
|
|
if err != nil {
|
|
t.Errorf("ParseBytes returned an error: %v", err)
|
|
}
|
|
|
|
if result != tc.expectedOutput {
|
|
t.Errorf("Expected %s, got %s", tc.expectedOutput, result)
|
|
}
|
|
})
|
|
}
|
|
})
|
|
|
|
// Test OpenRouterParser
|
|
t.Run("OpenRouterParser", func(t *testing.T) {
|
|
parser := NewOpenRouterParser(logger, config.LoadConfigOrDefault("config.toml"))
|
|
|
|
for _, tc := range testCases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
result, err := parser.ParseBytes([]byte(tc.responseJSON))
|
|
if err != nil {
|
|
t.Errorf("ParseBytes returned an error: %v", err)
|
|
}
|
|
|
|
if result != tc.expectedOutput {
|
|
t.Errorf("Expected %s, got %s", tc.expectedOutput, result)
|
|
}
|
|
})
|
|
}
|
|
})
|
|
}
|
|
|