completions_test.go

package blmstudio_test

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	blmstudio "github.com/bearaujus/blmstudio"
)

// ---------------------------------------------------------------------------
// helpers shared across all test files
// ---------------------------------------------------------------------------

func newTestClient(t *testing.T, srv *httptest.Server) blmstudio.Client {
	t.Helper()
	return blmstudio.New(blmstudio.WithBaseURL(srv.URL))
}

func writeJSON(w http.ResponseWriter, code int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	_ = json.NewEncoder(w).Encode(v)
}

func writeSSE(w http.ResponseWriter, lines []string) {
	w.Header().Set("Content-Type", "text/event-stream")
	for _, l := range lines {
		_, _ = w.Write([]byte(l + "\n"))
	}
}

// ---------------------------------------------------------------------------
// ChatCompletions
// ---------------------------------------------------------------------------

func TestChatCompletions_Happy(t *testing.T) {
	t.Parallel()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, 200, blmstudio.ChatCompletionResponse{
			ID:    "id1",
			Model: "test-model",
			Choices: []blmstudio.ChatCompletionChoice{{
				Message: blmstudio.Message{Role: blmstudio.RoleAssistant, Content: "hello"},
			}},
		})
	}))
	defer srv.Close()

	c := newTestClient(t, srv)
	resp, err := c.ChatCompletions(context.Background(), &blmstudio.ChatCompletionRequest{
		Model:    "test-model",
		Messages: []blmstudio.Message{{Role: blmstudio.RoleUser, Content: "hi"}},
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if resp.Choices[0].Message.Content != "hello" {
		t.Errorf("got %q, want %q", resp.Choices[0].Message.Content, "hello")
	}
}

func TestChatCompletions_EmptyPrompt(t *testing.T) {
	t.Parallel()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, 200, blmstudio.ChatCompletionResponse{})
	}))
	defer srv.Close()

	c := newTestClient(t, srv)
	_, err := c.ChatCompletions(context.Background(), &blmstudio.ChatCompletionRequest{
		Model:    "test-model",
		Messages: []blmstudio.Message{},
	})
	if !errors.Is(err, blmstudio.ErrEmptyPrompt) {
		t.Errorf("got %v, want ErrEmptyPrompt", err)
	}
}

func TestChatCompletions_4xx(t *testing.T) {
	t.Parallel()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, 400, map[string]any{"error": "bad request", "message": "invalid input"})
	}))
	defer srv.Close()

	c := newTestClient(t, srv)
	_, err := c.ChatCompletions(context.Background(), &blmstudio.ChatCompletionRequest{
		Model:    "test-model",
		Messages: []blmstudio.Message{{Role: blmstudio.RoleUser, Content: "hi"}},
	})
	if _, ok := blmstudio.IsAPIError(err); !ok {
		t.Errorf("expected APIError, got %v", err)
	}
}

func TestChatCompletions_5xx(t *testing.T) {
	t.Parallel()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(500)
	}))
	defer srv.Close()

	c := newTestClient(t, srv)
	_, err := c.ChatCompletions(context.Background(), &blmstudio.ChatCompletionRequest{
		Model:    "test-model",
		Messages: []blmstudio.Message{{Role: blmstudio.RoleUser, Content: "hi"}},
	})
	if err == nil {
		t.Fatal("expected error for 500, got nil")
	}
}

func TestChatCompletions_MalformedJSON(t *testing.T) {
	t.Parallel()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(200)
		_, _ = w.Write([]byte("{not json"))
	}))
	defer srv.Close()

	c := newTestClient(t, srv)
	_, err := c.ChatCompletions(context.Background(), &blmstudio.ChatCompletionRequest{
		Model:    "test-model",
		Messages: []blmstudio.Message{{Role: blmstudio.RoleUser, Content: "hi"}},
	})
	if err == nil {
		t.Fatal("expected error for malformed JSON, got nil")
	}
}

func TestChatCompletions_Timeout(t *testing.T) {
	t.Parallel()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(500 * time.Millisecond)
		writeJSON(w, 200, blmstudio.ChatCompletionResponse{})
	}))
	defer srv.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()

	c := newTestClient(t, srv)
	_, err := c.ChatCompletions(ctx, &blmstudio.ChatCompletionRequest{
		Model:    "test-model",
		Messages: []blmstudio.Message{{Role: blmstudio.RoleUser, Content: "hi"}},
	})
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Errorf("got %v, want DeadlineExceeded", err)
	}
}

func TestChatCompletionsStream_Happy(t *testing.T) {
	t.Parallel()
	chunks := []string{
		`data: {"id":"1","choices":[{"delta":{"content":"he"},"finish_reason":null}]}`,
		`data: {"id":"1","choices":[{"delta":{"content":"llo"},"finish_reason":null}]}`,
		`data: [DONE]`,
	}
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		writeSSE(w, chunks)
	}))
	defer srv.Close()

	c := newTestClient(t, srv)
	stream, err := c.ChatCompletionsStream(context.Background(), &blmstudio.ChatCompletionRequest{
		Model:    "test-model",
		Messages: []blmstudio.Message{{Role: blmstudio.RoleUser, Content: "hi"}},
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	defer stream.Close()

	var sb strings.Builder
	for stream.Next() {
		ev := stream.Event()
		if len(ev.Choices) > 0 {
			sb.WriteString(ev.Choices[0].Delta.Content)
		}
	}
	if err := stream.Err(); err != nil {
		t.Fatalf("stream error: %v", err)
	}
	if sb.String() != "hello" {
		t.Errorf("got %q, want %q", sb.String(), "hello")
	}
}

// ---------------------------------------------------------------------------
// Interface compliance
// ---------------------------------------------------------------------------

func TestClientInterface(t *testing.T) {
	t.Parallel()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, 200, map[string]any{})
	}))
	defer srv.Close()

	var _ blmstudio.Client = blmstudio.New(blmstudio.WithBaseURL(srv.URL))
}