69 lines
2.1 KiB
Go
69 lines
2.1 KiB
Go
package ai
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestFactory(t *testing.T) {
|
|
if _, e := New("bogus", "k"); e == nil || e.Error() != "AI_PROVIDER_INVALID" {
|
|
t.Fatalf("expected AI_PROVIDER_INVALID, got %v", e)
|
|
}
|
|
if _, e := New(ProviderSpark, ""); e == nil || e.Error() != "AI_KEY_MISSING" {
|
|
t.Fatalf("expected AI_KEY_MISSING, got %v", e)
|
|
}
|
|
p, e := New(ProviderDeepSeek, "key")
|
|
if e != nil || p.Name() != ProviderDeepSeek {
|
|
t.Fatalf("unexpected: %v %v", p, e)
|
|
}
|
|
p, e = New(ProviderSpark, "key")
|
|
if e != nil || p.Name() != ProviderSpark {
|
|
t.Fatalf("unexpected: %v %v", p, e)
|
|
}
|
|
}
|
|
|
|
func TestChatStreamSSE(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/chat/completions" {
|
|
t.Errorf("unexpected path %s", r.URL.Path)
|
|
}
|
|
if got := r.Header.Get("Authorization"); got != "Bearer test-key" {
|
|
t.Errorf("unexpected auth header %q", got)
|
|
}
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"你好\"}}]}\n\n"))
|
|
w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"，世界\"}}]}\n\n"))
|
|
w.Write([]byte("data: [DONE]\n\n"))
|
|
}))
|
|
defer srv.Close()
|
|
p := &openAICompatible{name: "test", baseURL: srv.URL, model: "m", apiKey: "test-key", client: srv.Client()}
|
|
stream, e := p.ChatStream(context.Background(), []Message{{Role: "user", Content: "hi"}})
|
|
if e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
var sb strings.Builder
|
|
for c := range stream {
|
|
if c.Err != nil {
|
|
t.Fatal(c.Err)
|
|
}
|
|
sb.WriteString(c.Content)
|
|
}
|
|
if sb.String() != "你好，世界" {
|
|
t.Fatalf("unexpected content %q", sb.String())
|
|
}
|
|
}
|
|
|
|
func TestChatStreamAuthError(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusUnauthorized)
|
|
}))
|
|
defer srv.Close()
|
|
p := &openAICompatible{name: "test", baseURL: srv.URL, model: "m", apiKey: "bad", client: srv.Client()}
|
|
if _, e := p.ChatStream(context.Background(), nil); e == nil || e.Error() != "AI_KEY_INVALID" {
|
|
t.Fatalf("expected AI_KEY_INVALID, got %v", e)
|
|
}
|
|
}
|