package vllm import ( "context" "encoding/json" "fmt" "net/http" "net/http/httptest" "strings" "sync" "testing" "go.uber.org/zap" "iop/apps/node/internal/runtime" "iop/packages/go/config" ) type fakeSink struct { mu sync.Mutex events []runtime.RuntimeEvent } func (s *fakeSink) Emit(_ context.Context, event runtime.RuntimeEvent) error { s.mu.Lock() defer s.mu.Unlock() s.events = append(s.events, event) return nil } func (s *fakeSink) all() []runtime.RuntimeEvent { s.mu.Lock() defer s.mu.Unlock() return append([]runtime.RuntimeEvent(nil), s.events...) } func TestVllmCapabilitiesQueryModels(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/v1/models" { t.Fatalf("unexpected path %s", r.URL.Path) } _, _ = w.Write([]byte(`{"object":"list","data":[{"id":"model-a"},{"id":"model-b"}]}`)) })) defer server.Close() adapter := New(config.VllmConf{ Endpoint: server.URL, Capacity: 10, MaxQueue: 20, QueueTimeoutMS: 5000, RequestTimeoutMS: 10000, }, zap.NewNop()) caps, err := adapter.Capabilities(context.Background()) if err != nil { t.Fatalf("Capabilities failed: %v", err) } if got := strings.Join(caps.Targets, ","); got != "model-a,model-b" { t.Fatalf("targets: got %q", got) } if caps.MaxConcurrency != 10 { t.Fatalf("expected MaxConcurrency 10, got %d", caps.MaxConcurrency) } if caps.MaxQueue != 20 { t.Fatalf("expected MaxQueue 20, got %d", caps.MaxQueue) } if caps.QueueTimeoutMS != 5000 { t.Fatalf("expected QueueTimeoutMS 5000, got %d", caps.QueueTimeoutMS) } if caps.RequestTimeoutMS != 10000 { t.Fatalf("expected RequestTimeoutMS 10000, got %d", caps.RequestTimeoutMS) } // Verify default capacity fallback adapterDefault := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop()) capsDefault, err := adapterDefault.Capabilities(context.Background()) if err != nil { t.Fatalf("Capabilities failed: %v", err) } if capsDefault.MaxConcurrency != 8 { t.Fatalf("expected default MaxConcurrency 8, got %d", capsDefault.MaxConcurrency) } } func TestVllmExecuteStreamsDeltas(t *testing.T) { var gotModel string var gotMessages int var gotStream bool server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/v1/chat/completions" { t.Fatalf("unexpected path %s", r.URL.Path) } var req struct { Model string `json:"model"` Messages []any `json:"messages"` Stream bool `json:"stream"` } if err := json.NewDecoder(r.Body).Decode(&req); err != nil { t.Fatalf("decode request: %v", err) } gotModel = req.Model gotMessages = len(req.Messages) gotStream = req.Stream w.Header().Set("Content-Type", "text/event-stream") _, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"hello "}}]}`) _, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"world"}}]}`) _, _ = fmt.Fprintf(w, "data: [DONE]\n\n") })) defer server.Close() adapter := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop()) sink := &fakeSink{} err := adapter.Execute(context.Background(), runtime.ExecutionSpec{ RunID: "run-1", Target: "llama-3", Input: map[string]any{"prompt": "say hello"}, }, sink) if err != nil { t.Fatalf("Execute failed: %v", err) } if gotModel != "llama-3" { t.Fatalf("model: got %q", gotModel) } if gotMessages != 1 { t.Fatalf("messages: got %d", gotMessages) } if !gotStream { t.Fatal("expected stream=true") } events := sink.all() if len(events) != 4 { t.Fatalf("expected 4 events (start+delta+delta+complete), got %d: %+v", len(events), events) } if events[0].Type != runtime.EventTypeStart { t.Fatalf("expected start event, got %+v", events[0]) } if events[1].Delta+events[2].Delta != "hello world" { t.Fatalf("unexpected deltas: %q + %q", events[1].Delta, events[2].Delta) } if events[3].Type != runtime.EventTypeComplete { t.Fatalf("expected complete event, got %+v", events[3]) } } func TestVllmExecuteUsesMessagesInput(t *testing.T) { var gotMessages []map[string]any server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { var req struct { Messages []map[string]any `json:"messages"` } if err := json.NewDecoder(r.Body).Decode(&req); err != nil { t.Fatalf("decode request: %v", err) } gotMessages = req.Messages w.Header().Set("Content-Type", "text/event-stream") _, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"ok"}}]}`) _, _ = fmt.Fprintf(w, "data: [DONE]\n\n") })) defer server.Close() adapter := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop()) sink := &fakeSink{} err := adapter.Execute(context.Background(), runtime.ExecutionSpec{ RunID: "run-2", Target: "llama-3", Input: map[string]any{ "messages": []any{ map[string]any{"role": "system", "content": "You are helpful."}, map[string]any{"role": "user", "content": "hi"}, }, }, }, sink) if err != nil { t.Fatalf("Execute failed: %v", err) } if len(gotMessages) != 2 { t.Fatalf("messages: got %d", len(gotMessages)) } if gotMessages[0]["role"] != "system" || gotMessages[1]["role"] != "user" { t.Fatalf("unexpected messages: %+v", gotMessages) } } func TestVllmExecuteEmitsErrorForHTTPFailure(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { http.Error(w, "service unavailable", http.StatusServiceUnavailable) })) defer server.Close() adapter := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop()) sink := &fakeSink{} err := adapter.Execute(context.Background(), runtime.ExecutionSpec{ RunID: "run-err", Target: "llama-3", Input: map[string]any{"prompt": "hi"}, }, sink) if err == nil { t.Fatal("expected error") } events := sink.all() if len(events) < 2 || events[1].Type != runtime.EventTypeError { t.Fatalf("expected error event after start, got %+v", events) } if !strings.Contains(events[1].Error, "503") { t.Fatalf("expected status in error, got %q", events[1].Error) } } func TestVllmInstanceKey(t *testing.T) { adapter := New(config.VllmConf{Endpoint: "http://localhost:8000"}, zap.NewNop(), "vllm-gpu") caps, _ := adapter.Capabilities(context.Background()) if caps.InstanceKey != "vllm-gpu" { t.Fatalf("InstanceKey: got %q, want vllm-gpu", caps.InstanceKey) } if caps.AdapterName != Name { t.Fatalf("AdapterName: got %q, want %q", caps.AdapterName, Name) } } func TestVllmProbeProviderAvailability(t *testing.T) { t.Run("200_ok_and_target_hit", func(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/v1/models" { t.Fatalf("unexpected path %s", r.URL.Path) } _, _ = w.Write([]byte(`{"object":"list","data":[{"id":"model-a"},{"id":"model-b"}]}`)) })) defer server.Close() adapter := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop()) res, err := adapter.ProbeProvider(context.Background(), "model-a") if err != nil { t.Fatalf("ProbeProvider failed: %v", err) } if res.Status != runtime.ProviderStatusAvailable { t.Errorf("expected Status available, got %s", res.Status) } if len(res.Targets) != 2 || res.Targets[0] != "model-a" { t.Errorf("unexpected Targets: %+v", res.Targets) } }) t.Run("200_ok_and_target_miss", func(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = w.Write([]byte(`{"object":"list","data":[{"id":"model-a"}]}`)) })) defer server.Close() adapter := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop()) res, err := adapter.ProbeProvider(context.Background(), "model-b") if err != nil { t.Fatalf("ProbeProvider failed: %v", err) } if res.Status != runtime.ProviderStatusUnavailable { t.Errorf("expected Status unavailable, got %s", res.Status) } if !strings.Contains(res.Detail, "not found") { t.Errorf("expected 'not found' in detail, got %s", res.Detail) } }) t.Run("500_internal_error", func(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusInternalServerError) })) defer server.Close() adapter := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop()) res, err := adapter.ProbeProvider(context.Background(), "model-a") if err != nil { t.Fatalf("ProbeProvider failed: %v", err) } if res.Status != runtime.ProviderStatusUnavailable { t.Errorf("expected Status unavailable, got %s", res.Status) } }) t.Run("empty_target", func(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = w.Write([]byte(`{"object":"list","data":[{"id":"model-a"}]}`)) })) defer server.Close() adapter := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop()) res, err := adapter.ProbeProvider(context.Background(), "") if err != nil { t.Fatalf("ProbeProvider failed: %v", err) } if res.Status != runtime.ProviderStatusAvailable { t.Errorf("expected Status available, got %s", res.Status) } }) }