87 lines
2.4 KiB
Go
87 lines
2.4 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestFakeLLMProvider_ChatCompletion(t *testing.T) {
|
|
provider := NewFakeLLMProvider()
|
|
|
|
resp, err := provider.ChatCompletion(context.Background(), ChatRequest{
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Hello world"},
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
assert.Contains(t, resp.Choices[0].Message.Content, "Hello world")
|
|
assert.Equal(t, "fake-ai-model", resp.Model)
|
|
}
|
|
|
|
func TestFakeLLMProvider_ChatCompletion_NoUserMessage(t *testing.T) {
|
|
provider := NewFakeLLMProvider()
|
|
|
|
resp, err := provider.ChatCompletion(context.Background(), ChatRequest{
|
|
Messages: []ChatMessage{
|
|
{Role: "system", Content: "You are helpful"},
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
assert.Contains(t, resp.Choices[0].Message.Content, "fake AI assistant")
|
|
}
|
|
|
|
func TestFakeLLMProvider_ChatCompletion_CallCount(t *testing.T) {
|
|
provider := NewFakeLLMProvider()
|
|
|
|
_, _ = provider.ChatCompletion(context.Background(), ChatRequest{
|
|
Messages: []ChatMessage{{Role: "user", Content: "msg1"}},
|
|
})
|
|
_, _ = provider.ChatCompletion(context.Background(), ChatRequest{
|
|
Messages: []ChatMessage{{Role: "user", Content: "msg2"}},
|
|
})
|
|
|
|
resp, err := provider.ChatCompletion(context.Background(), ChatRequest{
|
|
Messages: []ChatMessage{{Role: "user", Content: "msg3"}},
|
|
})
|
|
require.NoError(t, err)
|
|
assert.Contains(t, resp.Choices[0].Message.Content, "#3")
|
|
}
|
|
|
|
func TestIsFakeAIEnabled(t *testing.T) {
|
|
// Should return false when env var is not set
|
|
result := IsFakeAIEnabled()
|
|
assert.False(t, result)
|
|
}
|
|
|
|
func TestFakeLLMProvider_CreateEmbedding(t *testing.T) {
|
|
provider := NewFakeLLMProvider()
|
|
|
|
resp, err := provider.CreateEmbedding(context.Background(), EmbeddingRequest{
|
|
Input: []string{"test text"},
|
|
})
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, resp.Data)
|
|
assert.NotEmpty(t, resp.Data[0].Embedding)
|
|
}
|
|
|
|
func TestFakeLLMProvider_ChatCompletionStream(t *testing.T) {
|
|
provider := NewFakeLLMProvider()
|
|
|
|
var chunks []string
|
|
err := provider.ChatCompletionStream(context.Background(), ChatRequest{
|
|
Messages: []ChatMessage{{Role: "user", Content: "stream test"}},
|
|
}, func(chunk StreamChunk) error {
|
|
for _, choice := range chunk.Choices {
|
|
if choice.Delta.Content != "" {
|
|
chunks = append(chunks, choice.Delta.Content)
|
|
}
|
|
}
|
|
return nil
|
|
})
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, chunks)
|
|
}
|