Spaces:
Paused
Paused
File size: 4,232 Bytes
4ac9af9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | package executor
import (
"testing"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
"github.com/tidwall/gjson"
)
func TestEnsureImageGenerationTool_NoTools(t *testing.T) {
body := []byte(`{"model":"gpt-5.4","input":"draw a cat"}`)
result := ensureImageGenerationTool(body, "gpt-5.4", nil)
tools := gjson.GetBytes(result, "tools")
if !tools.IsArray() {
t.Fatalf("expected tools array, got %v", tools.Type)
}
arr := tools.Array()
if len(arr) != 1 {
t.Fatalf("expected 1 tool, got %d", len(arr))
}
if arr[0].Get("type").String() != "image_generation" {
t.Fatalf("expected type=image_generation, got %s", arr[0].Get("type").String())
}
if arr[0].Get("output_format").String() != "png" {
t.Fatalf("expected output_format=png, got %s", arr[0].Get("output_format").String())
}
}
func TestEnsureImageGenerationTool_ExistingToolsWithoutImageGen(t *testing.T) {
body := []byte(`{"model":"gpt-5.4","tools":[{"type":"function","name":"get_weather","parameters":{}}]}`)
result := ensureImageGenerationTool(body, "gpt-5.4", nil)
tools := gjson.GetBytes(result, "tools")
arr := tools.Array()
if len(arr) != 2 {
t.Fatalf("expected 2 tools, got %d", len(arr))
}
if arr[0].Get("type").String() != "function" {
t.Fatalf("expected first tool type=function, got %s", arr[0].Get("type").String())
}
if arr[1].Get("type").String() != "image_generation" {
t.Fatalf("expected second tool type=image_generation, got %s", arr[1].Get("type").String())
}
}
func TestEnsureImageGenerationTool_AlreadyPresent(t *testing.T) {
body := []byte(`{"model":"gpt-5.4","tools":[{"type":"image_generation","output_format":"webp"},{"type":"function","name":"f1"}]}`)
result := ensureImageGenerationTool(body, "gpt-5.4", nil)
tools := gjson.GetBytes(result, "tools")
arr := tools.Array()
if len(arr) != 2 {
t.Fatalf("expected 2 tools (no duplicate), got %d", len(arr))
}
if arr[0].Get("output_format").String() != "webp" {
t.Fatalf("expected original output_format=webp preserved, got %s", arr[0].Get("output_format").String())
}
}
func TestEnsureImageGenerationTool_EmptyToolsArray(t *testing.T) {
body := []byte(`{"model":"gpt-5.4","tools":[]}`)
result := ensureImageGenerationTool(body, "gpt-5.4", nil)
tools := gjson.GetBytes(result, "tools")
arr := tools.Array()
if len(arr) != 1 {
t.Fatalf("expected 1 tool, got %d", len(arr))
}
if arr[0].Get("type").String() != "image_generation" {
t.Fatalf("expected type=image_generation, got %s", arr[0].Get("type").String())
}
}
func TestEnsureImageGenerationTool_WebSearchAndImageGen(t *testing.T) {
body := []byte(`{"model":"gpt-5.4","tools":[{"type":"web_search"}]}`)
result := ensureImageGenerationTool(body, "gpt-5.4", nil)
tools := gjson.GetBytes(result, "tools")
arr := tools.Array()
if len(arr) != 2 {
t.Fatalf("expected 2 tools, got %d", len(arr))
}
if arr[0].Get("type").String() != "web_search" {
t.Fatalf("expected first tool type=web_search, got %s", arr[0].Get("type").String())
}
if arr[1].Get("type").String() != "image_generation" {
t.Fatalf("expected second tool type=image_generation, got %s", arr[1].Get("type").String())
}
}
func TestEnsureImageGenerationTool_GPT53CodexSparkDoesNotInjectTool(t *testing.T) {
body := []byte(`{"model":"gpt-5.3-codex-spark","input":"draw a cat"}`)
result := ensureImageGenerationTool(body, "gpt-5.3-codex-spark", nil)
if string(result) != string(body) {
t.Fatalf("expected body to be unchanged, got %s", string(result))
}
if gjson.GetBytes(result, "tools").Exists() {
t.Fatalf("expected no tools for gpt-5.3-codex-spark, got %s", gjson.GetBytes(result, "tools").Raw)
}
}
func TestEnsureImageGenerationTool_FreeCodexAuthDoesNotInjectTool(t *testing.T) {
body := []byte(`{"model":"gpt-5.4","input":"draw a cat"}`)
freeAuth := &cliproxyauth.Auth{
Provider: "codex",
Attributes: map[string]string{"plan_type": "free"},
}
result := ensureImageGenerationTool(body, "gpt-5.4", freeAuth)
if string(result) != string(body) {
t.Fatalf("expected body to be unchanged, got %s", string(result))
}
if gjson.GetBytes(result, "tools").Exists() {
t.Fatalf("expected no tools for free codex auth, got %s", gjson.GetBytes(result, "tools").Raw)
}
}
|