package openai import ( "context" "encoding/json" "net/http" "net/http/httptest" "strings" "testing" "time" edgeservice "iop/apps/edge/internal/service" "iop/packages/config" iop "iop/proto/gen/iop" ) type fakeRunService struct { req edgeservice.SubmitRunRequest events chan *iop.RunEvent } func (s *fakeRunService) SubmitRun(_ context.Context, req edgeservice.SubmitRunRequest) (*edgeservice.RunHandle, error) { s.req = req return &edgeservice.RunHandle{ RunDispatch: edgeservice.RunDispatch{ RunID: "run-test", Target: req.Target, TimeoutSec: 5, }, RunStream: edgeservice.RunStream{ Events: s.events, NodeEvents: make(chan *iop.EdgeNodeEvent), }, }, nil } func TestChatCompletionsDispatchesConfiguredOllamaTarget(t *testing.T) { fake := &fakeRunService{events: make(chan *iop.RunEvent, 3)} fake.events <- &iop.RunEvent{Type: "delta", Delta: "hello"} fake.events <- &iop.RunEvent{Type: "delta", Delta: " world"} fake.events <- &iop.RunEvent{Type: "complete", Usage: &iop.Usage{InputTokens: 2, OutputTokens: 2}} srv := NewServer(config.EdgeOpenAIConf{ Adapter: "ollama", Target: "llama-fixed", SessionID: "cline", TimeoutSec: 15, }, fake, nil) req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{ "model":"client-model", "messages":[{"role":"system","content":"brief"},{"role":"user","content":"say hello"}] }`)) w := httptest.NewRecorder() srv.handleChatCompletions(w, req) if w.Code != http.StatusOK { t.Fatalf("status: got %d body=%s", w.Code, w.Body.String()) } if fake.req.Adapter != "ollama" || fake.req.Target != "llama-fixed" { t.Fatalf("dispatch target mismatch: %+v", fake.req) } if fake.req.SessionID != "cline" || fake.req.TimeoutSec != 15 { t.Fatalf("execution config mismatch: %+v", fake.req) } if !strings.Contains(fake.req.Prompt, "system: brief") || !strings.Contains(fake.req.Prompt, "user: say hello") { t.Fatalf("prompt did not include messages: %q", fake.req.Prompt) } var resp chatCompletionResponse if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { t.Fatalf("decode response: %v", err) } if resp.Choices[0].Message.Content != "hello world" { t.Fatalf("content: got %q", resp.Choices[0].Message.Content) } if resp.Usage == nil || resp.Usage.TotalTokens != 4 { t.Fatalf("usage: %+v", resp.Usage) } } func TestChatCompletionsUsesRequestModelWhenNoConfiguredTarget(t *testing.T) { fake := &fakeRunService{events: make(chan *iop.RunEvent, 2)} fake.events <- &iop.RunEvent{Type: "delta", Delta: "ok"} fake.events <- &iop.RunEvent{Type: "complete"} srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama"}, fake, nil) req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{ "model":"from-request", "messages":[{"role":"user","content":"hi"}] }`)) w := httptest.NewRecorder() srv.handleChatCompletions(w, req) if w.Code != http.StatusOK { t.Fatalf("status: got %d body=%s", w.Code, w.Body.String()) } if fake.req.Target != "from-request" { t.Fatalf("target: got %q", fake.req.Target) } } func TestChatCompletionsStreamsSSE(t *testing.T) { fake := &fakeRunService{events: make(chan *iop.RunEvent, 3)} fake.events <- &iop.RunEvent{Type: "delta", Delta: "hi"} fake.events <- &iop.RunEvent{Type: "complete"} srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama"}, fake, nil) req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{ "model":"stream-model", "stream":true, "messages":[{"role":"user","content":"hi"}] }`)) w := httptest.NewRecorder() srv.handleChatCompletions(w, req) if w.Code != http.StatusOK { t.Fatalf("status: got %d body=%s", w.Code, w.Body.String()) } body := w.Body.String() if !strings.Contains(body, `"content":"hi"`) || !strings.Contains(body, "data: [DONE]") { t.Fatalf("unexpected SSE body:\n%s", body) } } func TestModelsUsesConfiguredModelsOrTarget(t *testing.T) { srv := NewServer(config.EdgeOpenAIConf{Target: "fallback-model"}, &fakeRunService{}, nil) req := httptest.NewRequest(http.MethodGet, "/v1/models", nil) w := httptest.NewRecorder() srv.handleModels(w, req) if w.Code != http.StatusOK { t.Fatalf("status: got %d", w.Code) } if !strings.Contains(w.Body.String(), "fallback-model") { t.Fatalf("expected target model, got %s", w.Body.String()) } srv = NewServer(config.EdgeOpenAIConf{Models: []string{"a", "b"}}, &fakeRunService{}, nil) w = httptest.NewRecorder() srv.handleModels(w, req) if !strings.Contains(w.Body.String(), `"id":"a"`) || !strings.Contains(w.Body.String(), `"id":"b"`) { t.Fatalf("expected configured models, got %s", w.Body.String()) } } func TestCollectRunTextTimesOut(t *testing.T) { handle := &edgeservice.RunHandle{ RunDispatch: edgeservice.RunDispatch{TimeoutSec: 1}, RunStream: edgeservice.RunStream{ Events: make(chan *iop.RunEvent), NodeEvents: make(chan *iop.EdgeNodeEvent), }, } ctx, cancel := context.WithTimeout(context.Background(), 1500*time.Millisecond) defer cancel() _, _, err := collectRunText(ctx, handle) if err == nil { t.Fatal("expected timeout error") } }