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