package ai import ( "context" "encoding/json" "io" "net/http" "strings" "testing" ) type roundTripFunc func(*http.Request) (*http.Response, error) func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) } func TestValidateEndpoint(t *testing.T) { if err := validateEndpoint("https://api.deepseek.com/v1", false); err != nil { t.Fatalf("https 应直接放行: %v", err) } if err := validateEndpoint("http://127.0.0.1:11434/v1", false); err == nil { t.Fatal("未开启 allowInsecure 时 http 应拒绝") } if err := validateEndpoint("http://127.0.0.1:11434/v1", true); err != nil { t.Fatalf("allowInsecure 下本地 http 应放行: %v", err) } if err := validateEndpoint("ftp://example.com", true); err == nil { t.Fatal("非 HTTP/HTTPS scheme 应拒绝") } if err := validateEndpoint("not a url", false); err == nil { t.Fatal("非法 URL 应拒绝") } } func TestEndpointJoin(t *testing.T) { if got := endpoint("https://api.example.com/v1/", "chat/completions"); got != "https://api.example.com/v1/chat/completions" { t.Fatalf("拼接错误: %s", got) } if got := endpoint("https://api.example.com/v1/responses", "responses"); got != "https://api.example.com/v1/responses" { t.Fatalf("已含 suffix 不应重复: %s", got) } if got := anthropicEndpoint("https://api.anthropic.com"); got != "https://api.anthropic.com/v1/messages" { t.Fatalf("anthropic 无 /v1 应补: %s", got) } if got := anthropicEndpoint("https://api.anthropic.com/v1"); got != "https://api.anthropic.com/v1/messages" { t.Fatalf("anthropic 已含 /v1 不应重复: %s", got) } } func TestNewRemoteProviderInfersProtocolAndAuth(t *testing.T) { legacy := NewRemoteProvider(Config{Provider: "openai", BaseURL: "https://api.openai.com/v1"}) if legacy.protocol != ProtoChat || legacy.authStyle != AuthBearer { t.Fatalf("旧配置应推断 openai-chat/bearer, got %s/%s", legacy.protocol, legacy.authStyle) } anthropic := NewRemoteProvider(Config{Provider: "anthropic", BaseURL: "https://api.anthropic.com"}) if anthropic.protocol != ProtoAnthropic || anthropic.authStyle != AuthXAPIKey { t.Fatalf("anthropic 预设应推断原生协议/x-api-key, got %s/%s", anthropic.protocol, anthropic.authStyle) } local := NewRemoteProvider(Config{Provider: "ollama", Protocol: ProtoChat, AuthStyle: AuthNone, BaseURL: "http://127.0.0.1:11434/v1", AllowInsecureHTTP: true}) if local.authStyle != AuthNone { t.Fatalf("本地免鉴权应保留 none, got %s", local.authStyle) } } func TestParseContentByProtocol(t *testing.T) { chatRaw, _ := json.Marshal(map[string]any{ "choices": []map[string]any{{"message": map[string]string{"role": "assistant", "content": " ```sql\nSELECT 1;\n``` "}}}, }) got, err := parseContent(ProtoChat, chatRaw, http.StatusOK) if err != nil || strings.TrimSpace(got) != "```sql\nSELECT 1;\n```" { t.Fatalf("chat 解析失败: %q, %v", got, err) } if s := extractSQL(got); s != "SELECT 1;" { t.Fatalf("sql 代码块提取错误: %q", s) } anthropicRaw, _ := json.Marshal(map[string]any{ "content": []map[string]any{{"type": "text", "text": "你好"}, {"type": "tool_use", "text": "忽略"}}, }) got, err = parseContent(ProtoAnthropic, anthropicRaw, http.StatusOK) if err != nil || got != "你好" { t.Fatalf("anthropic 解析失败: %q, %v", got, err) } responsesRaw, _ := json.Marshal(map[string]any{ "output": []map[string]any{{"content": []map[string]string{{"text": "SELECT "}, {"text": "2;"}}}}, }) got, err = parseContent(ProtoResponses, responsesRaw, http.StatusOK) if err != nil || got != "SELECT 2;" { t.Fatalf("responses 解析失败: %q, %v", got, err) } errRaw, _ := json.Marshal(map[string]any{"error": map[string]string{"message": "invalid api key"}}) if _, err := parseContent(ProtoChat, errRaw, http.StatusUnauthorized); err == nil || !strings.Contains(err.Error(), "invalid api key") { t.Fatalf("非 2xx 应提取服务端错误信息: %v", err) } } func TestBuildMessagesCoversActions(t *testing.T) { msgs := BuildMessages(ChatReq{Action: "optimize", SQL: "SELECT * FROM t", Driver: "postgresql", Database: "db"}) if len(msgs) != 2 || msgs[0].Role != "system" || msgs[1].Role != "user" { t.Fatalf("应返回 [system, user], got %+v", msgs) } if !strings.Contains(msgs[1].Content, "SELECT * FROM t") { t.Fatal("user 消息应包含 SQL") } } func TestAskChatsWithoutClaimingDatabaseExecution(t *testing.T) { request := ChatReq{Action: "ask", Prompt: "这个库里有哪些表?", Driver: "mysql", Database: "demo"} msgs := BuildMessages(request) if !strings.Contains(msgs[0].Content, "demo(mysql)") || !strings.Contains(msgs[0].Content, "不要声称已经查询") { t.Fatalf("Ask should carry context without pretending to query: %s", msgs[0].Content) } if strings.Contains(msgs[0].Content, "只输出可执行 SQL") { t.Fatal("Ask must not force a SQL response") } generate := BuildMessages(ChatReq{Action: "generate", Prompt: "查 users", Driver: "mysql", Database: "demo"}) if !strings.Contains(generate[0].Content, "生成一条 mysql 方言的 SQL") { t.Fatalf("SQL generation lost its dedicated instruction: %s", generate[0].Content) } if !strings.Contains(generate[0].Content, "澄清") || strings.Contains(generate[0].Content, "只输出可执行 SQL") { t.Fatalf("generate must allow clarification and not force executable-only: %s", generate[0].Content) } if !strings.Contains(generate[0].Content, "禁止编造") || !strings.Contains(generate[0].Content, "UPDATE_TIME") { t.Fatalf("system rules missing invent/UPDATE_TIME guards: %s", generate[0].Content) } if !strings.Contains(generate[0].Content, "换表") || !strings.Contains(generate[0].Content, "created_at") { t.Fatalf("follow-up no-invent column rules missing: %s", generate[0].Content) } offline := chatStub(request) if offline.SQL != "" || !strings.Contains(offline.Content, "离线体验无法回答") { t.Fatalf("offline Ask must not pretend to generate or query: %+v", offline) } } func TestAskRemoteRequestIsChatAndReturnsNoUnrequestedSQL(t *testing.T) { p := NewRemoteProvider(Config{Provider: "deepseek", Protocol: ProtoChat, BaseURL: "https://api.example.test/v1", APIKey: "test-secret", Model: "test-model"}) p.hc = &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { if req.URL.Path != "/v1/chat/completions" { t.Fatalf("unexpected endpoint: %s", req.URL) } if req.Header.Get("Authorization") != "Bearer test-secret" { t.Fatal("missing auth") } var payload struct { Messages []Message `json:"messages"` Tools []any `json:"tools"` } if err := json.NewDecoder(req.Body).Decode(&payload); err != nil { t.Fatal(err) } if len(payload.Messages) != 2 || len(payload.Tools) != 0 || !strings.Contains(payload.Messages[0].Content, "demo(mysql)") || strings.Contains(payload.Messages[0].Content, "只输出可执行 SQL") { t.Fatalf("Ask should be contextual chat without tools: %+v", payload) } return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(`{"choices":[{"message":{"content":"当前表结构上下文不足,无法确认表名。"}}]}`)), Header: make(http.Header)}, nil })} res, err := p.Complete(context.Background(), ChatReq{Action: "ask", Prompt: "有哪些表?", Driver: "mysql", Database: "demo"}) if err != nil || res.SQL != "" || !strings.Contains(res.Content, "无法确认表名") { t.Fatalf("Ask must return a plain answer without invented SQL: %+v err=%v", res, err) } }