opsxdb/server/internal/ai/ai_test.go

package ai

import (
	"encoding/json"
	"net/http"
	"strings"
	"testing"
)

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")
	}
}