package llm import ( "context" "encoding/json" "errors" "fmt" "net/http" "net/http/httptest" "strings" "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) // --- Interface compliance test --- func TestOpenAIProvider_ImplementsProvider(t *testing.T) { var _ Provider = (*OpenAIProvider)(nil) } // --- Constructor tests --- func TestNewOpenAIProvider_Defaults(t *testing.T) { p := NewOpenAIProvider(OpenAIProviderConfig{ APIKey: "test-key", }) assert.Equal(t, "test-key", p.apiKey) assert.Equal(t, "https://api.openai.com/v1", p.baseURL) assert.Equal(t, "gpt-4", p.model) assert.Equal(t, "text-embedding-3-small", p.embedModel) assert.Equal(t, 3, p.maxRetries) } func TestNewOpenAIProvider_CustomConfig(t *testing.T) { p := NewOpenAIProvider(OpenAIProviderConfig{ APIKey: "custom-key", BaseURL: "https://ark.cn-beijing.volces.com/api/v3", Model: "doubao-pro-4k", EmbedModel: "doubao-embedding", MaxRetries: 5, Timeout: 120, }) assert.Equal(t, "custom-key", p.apiKey) assert.Equal(t, "https://ark.cn-beijing.volces.com/api/v3", p.baseURL) assert.Equal(t, "doubao-pro-4k", p.model) assert.Equal(t, "doubao-embedding", p.embedModel) assert.Equal(t, 5, p.maxRetries) } func TestNewOpenAIProvider_TrailingSlashTrimmed(t *testing.T) { p := NewOpenAIProvider(OpenAIProviderConfig{ APIKey: "key", BaseURL: "https://api.openai.com/v1/", }) assert.Equal(t, "https://api.openai.com/v1", p.baseURL) } // --- Request serialization tests --- func TestChatRequest_MarshalJSON(t *testing.T) { req := ChatRequest{ Model: "gpt-4", Messages: []ChatMessage{ {Role: "system", Content: "You are a helpful assistant."}, {Role: "user", Content: "Hello!"}, }, Temperature: 0.7, MaxTokens: 100, } body, err := json.Marshal(req) require.NoError(t, err) // Verify key fields are present assert.Contains(t, string(body), `"model":"gpt-4"`) assert.Contains(t, string(body), `"messages"`) assert.Contains(t, string(body), `"role":"system"`) assert.Contains(t, string(body), `"role":"user"`) assert.Contains(t, string(body), `"temperature":0.7`) assert.Contains(t, string(body), `"max_tokens":100`) } func TestChatRequest_MarshalJSON_WithTools(t *testing.T) { req := ChatRequest{ Model: "gpt-4", Messages: []ChatMessage{{Role: "user", Content: "Search for docs"}}, Tools: []ToolDefinition{ { Type: "function", Function: ToolFunction{ Name: "search_documents", Description: "Search knowledge base", Parameters: map[string]interface{}{ "type": "object", "properties": map[string]interface{}{ "query": map[string]interface{}{ "type": "string", }, }, }, }, }, }, } body, err := json.Marshal(req) require.NoError(t, err) assert.Contains(t, string(body), `"tools"`) assert.Contains(t, string(body), `"search_documents"`) } func TestEmbeddingRequest_MarshalJSON(t *testing.T) { req := EmbeddingRequest{ Model: "text-embedding-3-small", Input: []string{"Hello world", "Test embedding"}, } body, err := json.Marshal(req) require.NoError(t, err) assert.Contains(t, string(body), `"model":"text-embedding-3-small"`) assert.Contains(t, string(body), `"Hello world"`) assert.Contains(t, string(body), `"Test embedding"`) } func TestChatRequest_OmitEmptyFields(t *testing.T) { req := ChatRequest{ Model: "gpt-4", Messages: []ChatMessage{{Role: "user", Content: "Hi"}}, } body, err := json.Marshal(req) require.NoError(t, err) // temperature=0 should NOT be omitted since 0 is the zero value for float64 // but "omitempty" will omit it — this is expected behavior decoded := ChatRequest{} err = json.Unmarshal(body, &decoded) require.NoError(t, err) assert.Equal(t, "gpt-4", decoded.Model) assert.Equal(t, 1, len(decoded.Messages)) } // --- Response deserialization tests --- func TestChatResponse_UnmarshalJSON(t *testing.T) { raw := `{ "id": "chatcmpl-123", "object": "chat.completion", "created": 1677652288, "model": "gpt-4", "choices": [{ "index": 0, "message": {"role": "assistant", "content": "Hello! How can I help you?"}, "finish_reason": "stop" }], "usage": { "prompt_tokens": 9, "completion_tokens": 12, "total_tokens": 21 } }` var resp ChatResponse err := json.Unmarshal([]byte(raw), &resp) require.NoError(t, err) assert.Equal(t, "chatcmpl-123", resp.ID) assert.Equal(t, 1, len(resp.Choices)) assert.Equal(t, "assistant", resp.Choices[0].Message.Role) assert.Equal(t, "Hello! How can I help you?", resp.Choices[0].Message.Content) assert.Equal(t, "stop", resp.Choices[0].FinishReason) assert.Equal(t, 21, resp.Usage.TotalTokens) } func TestEmbeddingResponse_UnmarshalJSON(t *testing.T) { raw := `{ "object": "list", "data": [{ "object": "embedding", "index": 0, "embedding": [0.0023064255, -0.009327292, 0.015871] }], "model": "text-embedding-3-small", "usage": { "prompt_tokens": 5, "total_tokens": 5 } }` var resp EmbeddingResponse err := json.Unmarshal([]byte(raw), &resp) require.NoError(t, err) assert.Equal(t, 1, len(resp.Data)) assert.Equal(t, 0, resp.Data[0].Index) assert.Equal(t, 3, len(resp.Data[0].Embedding)) assert.InDelta(t, 0.0023064255, resp.Data[0].Embedding[0], 0.0001) } // --- Mock server tests --- func TestOpenAIProvider_ChatCompletion_MockServer(t *testing.T) { mockResp := ChatResponse{ ID: "chatcmpl-test", Object: "chat.completion", Model: "gpt-4", Choices: []ChatChoice{ { Index: 0, Message: ChatMessage{Role: "assistant", Content: "Mocked response"}, FinishReason: "stop", }, }, Usage: TokenUsage{PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15}, } server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { assert.Equal(t, "POST", r.Method) assert.Equal(t, "/chat/completions", r.URL.Path) assert.Equal(t, "Bearer test-key", r.Header.Get("Authorization")) assert.Equal(t, "application/json", r.Header.Get("Content-Type")) body, _ := json.Marshal(mockResp) w.WriteHeader(http.StatusOK) w.Write(body) })) defer server.Close() p := NewOpenAIProvider(OpenAIProviderConfig{ APIKey: "test-key", BaseURL: server.URL, Model: "gpt-4", MaxRetries: 0, }) req := ChatRequest{ Messages: []ChatMessage{{Role: "user", Content: "Hello"}}, } resp, err := p.ChatCompletion(context.Background(), req) require.NoError(t, err) assert.Equal(t, "chatcmpl-test", resp.ID) assert.Equal(t, "Mocked response", resp.Choices[0].Message.Content) } func TestOpenAIProvider_CreateEmbedding_MockServer(t *testing.T) { mockResp := EmbeddingResponse{ Object: "list", Data: []EmbeddingData{ { Object: "embedding", Index: 0, Embedding: []float64{0.1, 0.2, 0.3}, }, }, Model: "text-embedding-3-small", Usage: TokenUsage{PromptTokens: 3, TotalTokens: 3}, } server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { assert.Equal(t, "POST", r.Method) assert.Equal(t, "/embeddings", r.URL.Path) assert.Equal(t, "Bearer test-key", r.Header.Get("Authorization")) body, _ := json.Marshal(mockResp) w.WriteHeader(http.StatusOK) w.Write(body) })) defer server.Close() p := NewOpenAIProvider(OpenAIProviderConfig{ APIKey: "test-key", BaseURL: server.URL, MaxRetries: 0, }) req := EmbeddingRequest{ Input: []string{"Hello world"}, } resp, err := p.CreateEmbedding(context.Background(), req) require.NoError(t, err) assert.Equal(t, 1, len(resp.Data)) assert.InDelta(t, 0.1, resp.Data[0].Embedding[0], 0.001) } func TestOpenAIProvider_ChatCompletion_APIError(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusUnauthorized) w.Write([]byte(`{"error":{"message":"Invalid API key","type":"invalid_request_error","code":"invalid_api_key"}}`)) })) defer server.Close() p := NewOpenAIProvider(OpenAIProviderConfig{ APIKey: "bad-key", BaseURL: server.URL, MaxRetries: 0, }) req := ChatRequest{ Messages: []ChatMessage{{Role: "user", Content: "Hello"}}, } resp, err := p.ChatCompletion(context.Background(), req) assert.Nil(t, resp) require.Error(t, err) var apiErr *APIError require.True(t, errors.As(err, &apiErr), "error should wrap APIError") assert.Equal(t, http.StatusUnauthorized, apiErr.StatusCode) assert.Equal(t, "Invalid API key", apiErr.Message) assert.Equal(t, "invalid_request_error", apiErr.Type) } func TestOpenAIProvider_ChatCompletion_RetryOn5xx(t *testing.T) { callCount := 0 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { callCount++ if callCount < 3 { w.WriteHeader(http.StatusInternalServerError) w.Write([]byte(`{"error":{"message":"Internal server error"}}`)) return } mockResp := ChatResponse{ ID: "chatcmpl-retry", Choices: []ChatChoice{{Message: ChatMessage{Role: "assistant", Content: "Success after retry"}}}, } body, _ := json.Marshal(mockResp) w.WriteHeader(http.StatusOK) w.Write(body) })) defer server.Close() p := NewOpenAIProvider(OpenAIProviderConfig{ APIKey: "test-key", BaseURL: server.URL, MaxRetries: 3, Timeout: 5, }) req := ChatRequest{ Messages: []ChatMessage{{Role: "user", Content: "Hello"}}, } resp, err := p.ChatCompletion(context.Background(), req) require.NoError(t, err) assert.Equal(t, "Success after retry", resp.Choices[0].Message.Content) assert.Equal(t, 3, callCount) } func TestOpenAIProvider_ChatCompletion_NoRetryOn4xx(t *testing.T) { callCount := 0 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { callCount++ w.WriteHeader(http.StatusBadRequest) w.Write([]byte(`{"error":{"message":"Bad request"}}`)) })) defer server.Close() p := NewOpenAIProvider(OpenAIProviderConfig{ APIKey: "test-key", BaseURL: server.URL, MaxRetries: 3, Timeout: 5, }) req := ChatRequest{ Messages: []ChatMessage{{Role: "user", Content: "Hello"}}, } resp, err := p.ChatCompletion(context.Background(), req) assert.Nil(t, resp) require.Error(t, err) assert.Equal(t, 1, callCount, "should not retry on 400") } func TestOpenAIProvider_DefaultModelApplied(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { var reqBody ChatRequest json.NewDecoder(r.Body).Decode(&reqBody) assert.Equal(t, "gpt-4", reqBody.Model, "default model should be applied when not specified") mockResp := ChatResponse{ID: "test"} body, _ := json.Marshal(mockResp) w.WriteHeader(http.StatusOK) w.Write(body) })) defer server.Close() p := NewOpenAIProvider(OpenAIProviderConfig{ APIKey: "key", BaseURL: server.URL, Model: "gpt-4", }) req := ChatRequest{ Messages: []ChatMessage{{Role: "user", Content: "Hello"}}, // Model intentionally left empty } _, err := p.ChatCompletion(context.Background(), req) require.NoError(t, err) } // --- SSE stream parsing test --- func TestOpenAIProvider_ParseSSEStream(t *testing.T) { sseData := `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":""}]} data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":""}]} data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":" world"},"finish_reason":""}]} data: [DONE] ` p := NewOpenAIProvider(OpenAIProviderConfig{APIKey: "test"}) var chunks []StreamChunk err := p.parseSSEStream(strings.NewReader(sseData), func(chunk StreamChunk) error { chunks = append(chunks, chunk) return nil }) require.NoError(t, err) assert.Equal(t, 3, len(chunks)) assert.Equal(t, "assistant", chunks[0].Choices[0].Delta.Role) assert.Equal(t, "Hello", chunks[1].Choices[0].Delta.Content) assert.Equal(t, " world", chunks[2].Choices[0].Delta.Content) } // --- Utility tests --- func TestParseFloatEmbedding(t *testing.T) { raw := []interface{}{float64(0.1), float32(0.2), int(3), int64(4), "5.5"} result := ParseFloatEmbedding(raw) assert.Equal(t, 5, len(result)) assert.InDelta(t, 0.1, result[0], 0.001) assert.InDelta(t, 0.2, result[1], 0.001) assert.InDelta(t, 3.0, result[2], 0.001) assert.InDelta(t, 4.0, result[3], 0.001) assert.InDelta(t, 5.5, result[4], 0.001) } // --- APIError tests --- func TestAPIError_Error(t *testing.T) { err := &APIError{ StatusCode: 401, Message: "Invalid API key", Type: "invalid_request_error", Code: "invalid_api_key", } assert.Equal(t, "API error (status 401): Invalid API key", err.Error()) } func TestIsNonRetriableError(t *testing.T) { // 4xx (except 429) should not retry assert.True(t, isNonRetriableError(&APIError{StatusCode: 400})) assert.True(t, isNonRetriableError(&APIError{StatusCode: 401})) assert.True(t, isNonRetriableError(&APIError{StatusCode: 403})) assert.True(t, isNonRetriableError(&APIError{StatusCode: 404})) // 429 (rate limit) should retry assert.False(t, isNonRetriableError(&APIError{StatusCode: 429})) // 5xx should retry assert.False(t, isNonRetriableError(&APIError{StatusCode: 500})) assert.False(t, isNonRetriableError(&APIError{StatusCode: 502})) assert.False(t, isNonRetriableError(&APIError{StatusCode: 503})) // Non-APIError should not be treated as non-retriable assert.False(t, isNonRetriableError(fmt.Errorf("some error"))) }