GovAI/server/pkg/llm/local_test.go

package llm

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

// 用一个 OpenAI 兼容的 mock 服务模拟本地 vLLM/Ollama，验证：
// 1) 本地 provider 的流式响应能被 TransformOpenAIStream 正确解析；
// 2) 未配置密钥时不发送 Authorization 头（本地无鉴权端点）。
func TestLocalProvider_StreamingAndNoAuthHeader(t *testing.T) {
	var gotAuth string
	var sawAuthHeader bool
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotAuth = r.Header.Get("Authorization")
		_, sawAuthHeader = r.Header["Authorization"]
		w.Header().Set("Content-Type", "text/event-stream")
		flusher, _ := w.(http.Flusher)
		for _, chunk := range []string{
			`{"id":"cmpl-1","model":"local-model","choices":[{"delta":{"content":"你好"}}]}`,
			`{"id":"cmpl-1","model":"local-model","choices":[{"delta":{"content":"，世界"}}]}`,
		} {
			fmt.Fprintf(w, "data: %s\n\n", chunk)
			if flusher != nil {
				flusher.Flush()
			}
		}
		fmt.Fprint(w, "data: [DONE]\n\n")
		if flusher != nil {
			flusher.Flush()
		}
	}))
	defer srv.Close()

	mgr := NewManager()
	// 密钥留空，模拟本地无鉴权端点
	mgr.Register("local", NewOpenAIProvider("", srv.URL, "local-model"))

	body, err := mgr.ChatStream(context.Background(), "local", &ChatRequest{
		Messages: []Message{{Role: RoleUser, Content: "hi"}},
	})
	if err != nil {
		t.Fatalf("ChatStream 出错: %v", err)
	}
	defer body.Close()

	var sb strings.Builder
	var ended bool
	if err := TransformOpenAIStream(body, func(ev StreamEvent) {
		if ev.Answer != "" {
			sb.WriteString(ev.Answer)
		}
		if ev.Event == "message_end" {
			ended = true
		}
	}); err != nil {
		t.Fatalf("解析流出错: %v", err)
	}

	if sb.String() != "你好，世界" {
		t.Fatalf("流式拼接结果不符: %q", sb.String())
	}
	if !ended {
		t.Fatal("未收到 message_end 事件")
	}
	if sawAuthHeader || gotAuth != "" {
		t.Fatalf("空密钥时不应发送 Authorization 头，实际: %q", gotAuth)
	}
}

// 配置了密钥时应发送 Authorization 头（云端/带鉴权的本地服务）。
func TestOpenAIProvider_SendsAuthHeaderWhenKeySet(t *testing.T) {
	var gotAuth string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotAuth = r.Header.Get("Authorization")
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{"id":"1","model":"m","choices":[{"message":{"content":"ok"}}],"usage":{"total_tokens":3}}`)
	}))
	defer srv.Close()

	p := NewOpenAIProvider("test-key", srv.URL, "m")
	resp, err := p.ChatCompletion(context.Background(), &ChatRequest{
		Messages: []Message{{Role: RoleUser, Content: "hi"}},
	})
	if err != nil {
		t.Fatalf("ChatCompletion 出错: %v", err)
	}
	if resp.Content != "ok" {
		t.Fatalf("响应内容不符: %q", resp.Content)
	}
	if gotAuth != "Bearer test-key" {
		t.Fatalf("应发送 Bearer 密钥头，实际: %q", gotAuth)
	}
}

// 未注册的 provider 名应回退到 fallback。
func TestManager_FallbackResolution(t *testing.T) {
	mgr := NewManager()
	mgr.Register("local", NewOpenAIProvider("", "http://localhost:9", "m"))
	mgr.SetFallback("local")

	if _, err := mgr.GetProvider("does-not-exist"); err != nil {
		t.Fatalf("未知 provider 应回退到 fallback，却报错: %v", err)
	}
}