grokkit/internal/grok/client_test.go

package grok

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"
)

func TestCleanCodeResponse(t *testing.T) {
	tests := []struct {
		name     string
		input    string
		expected string
	}{
		{
			name:     "removes markdown fences",
			input:    "```go\nfunc main() {}\n```",
			expected: "func main() {}",
		},
		{
			name:     "removes language tag",
			input:    "```python\nprint('hello')\n```",
			expected: "print('hello')",
		},
		{
			name:     "handles no fences",
			input:    "func main() {}",
			expected: "func main() {}",
		},
		{
			name:     "preserves internal blank lines",
			input:    "```\nline1\n\nline2\n```",
			expected: "line1\n\nline2",
		},
		{
			name:     "trims whitespace",
			input:    "  \n```\ncode\n```\n  ",
			expected: "code",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := CleanCodeResponse(tt.input)
			if result != tt.expected {
				t.Errorf("CleanCodeResponse() = %q, want %q", result, tt.expected)
			}
		})
	}
}

func sseServer(chunks []string) *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		for _, c := range chunks {
			fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":%q}}]}\n\n", c)
		}
		fmt.Fprintf(w, "data: [DONE]\n\n")
	}))
}

func TestStreamSilent(t *testing.T) {
	srv := sseServer([]string{"Hello", " ", "World"})
	defer srv.Close()

	client := &Client{APIKey: "test-key", BaseURL: srv.URL}
	got := client.StreamSilent([]map[string]string{{"role": "user", "content": "hi"}}, "test-model")
	if got != "Hello World" {
		t.Errorf("StreamSilent() = %q, want %q", got, "Hello World")
	}
}

func TestStream(t *testing.T) {
	srv := sseServer([]string{"foo", "bar"})
	defer srv.Close()

	client := &Client{APIKey: "test-key", BaseURL: srv.URL}
	got := client.Stream([]map[string]string{{"role": "user", "content": "hi"}}, "test-model")
	if got != "foobar" {
		t.Errorf("Stream() = %q, want %q", got, "foobar")
	}
}

func TestStreamWithTemp(t *testing.T) {
	srv := sseServer([]string{"response"})
	defer srv.Close()

	client := &Client{APIKey: "test-key", BaseURL: srv.URL}
	got := client.StreamWithTemp([]map[string]string{{"role": "user", "content": "hi"}}, "test-model", 0.5)
	if got != "response" {
		t.Errorf("StreamWithTemp() = %q, want %q", got, "response")
	}
}

func TestStreamDoneSignal(t *testing.T) {
	// Verifies that [DONE] stops processing and non-content chunks are skipped
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"}}]}\n\n")
		fmt.Fprintf(w, "data: [DONE]\n\n")
		// This line should never be processed
		fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"extra\"}}]}\n\n")
	}))
	defer srv.Close()

	client := &Client{APIKey: "test-key", BaseURL: srv.URL}
	got := client.StreamSilent(nil, "test-model")
	if got != "ok" {
		t.Errorf("got %q, want %q", got, "ok")
	}
}

func TestStreamEmptyResponse(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		fmt.Fprintf(w, "data: [DONE]\n\n")
	}))
	defer srv.Close()

	client := &Client{APIKey: "test-key", BaseURL: srv.URL}
	got := client.StreamSilent(nil, "test-model")
	if got != "" {
		t.Errorf("got %q, want empty string", got)
	}
}

func TestNewClient(t *testing.T) {
	// Save and restore env
	oldKey := os.Getenv("XAI_API_KEY")
	defer func() {
		if oldKey != "" {
			os.Setenv("XAI_API_KEY", oldKey)
		} else {
			os.Unsetenv("XAI_API_KEY")
		}
	}()

	os.Setenv("XAI_API_KEY", "test-key")
	client := NewClient()

	if client.APIKey != "test-key" {
		t.Errorf("NewClient() APIKey = %q, want %q", client.APIKey, "test-key")
	}

	if client.BaseURL != "https://api.x.ai/v1" {
		t.Errorf("NewClient() BaseURL = %q, want %q", client.BaseURL, "https://api.x.ai/v1")
	}
}