67 lines
1.8 KiB
Go
67 lines
1.8 KiB
Go
package utils_vllm
|
|
|
|
import (
|
|
"ai_scheduler/internal/config"
|
|
"context"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"github.com/cloudwego/eino/schema"
|
|
)
|
|
|
|
func newMockServer() *httptest.Server {
|
|
h := http.NewServeMux()
|
|
h.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte(`{"id":"cmpl-1","object":"chat.completion","created":173,"model":"x","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`))
|
|
})
|
|
return httptest.NewServer(h)
|
|
}
|
|
|
|
func Test_Vllm_Chat_Generate(t *testing.T) {
|
|
cfg, err := config.LoadConfig("../../../config/config_test.yaml")
|
|
if err != nil {
|
|
t.Fatalf("load config: %v", err)
|
|
}
|
|
|
|
ctx := context.Background()
|
|
client, _, err := NewClient(cfg)
|
|
if err != nil {
|
|
t.Fatalf("new client: %v", err)
|
|
}
|
|
|
|
msgs := []*schema.Message{{Role: schema.User, Content: "hi"}}
|
|
out, err := client.Chat(ctx, msgs)
|
|
if err != nil {
|
|
t.Fatalf("chat generate: %v", err)
|
|
}
|
|
if out == nil || out.Content != "ok" {
|
|
t.Fatalf("unexpected content: %v", out)
|
|
}
|
|
|
|
t.Logf("结果: %v", out)
|
|
}
|
|
|
|
func Test_Vllm_RecognizeWithImg(t *testing.T) {
|
|
cfg, err := config.LoadConfig("../../../config/config_test.yaml")
|
|
if err != nil {
|
|
t.Fatalf("load config: %v", err)
|
|
}
|
|
|
|
ctx := context.Background()
|
|
client, _, err := NewClient(cfg)
|
|
if err != nil {
|
|
t.Fatalf("new client: %v", err)
|
|
}
|
|
|
|
out, err := client.RecognizeWithImg(ctx, "sys", "user", []string{"https://img0.baidu.com/it/u=910428455,194434251&fm=253&app=138&f=JPEG?w=1122&h=800"})
|
|
if err != nil {
|
|
t.Fatalf("recognize with img: %v", err)
|
|
}
|
|
if out == nil || out.Content != "ok" {
|
|
t.Fatalf("unexpected content: %v", out)
|
|
}
|
|
}
|