grokkit/internal/grok/client_test.go

package grok

import (
	"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 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")
	}
}