package llm
import (
"testing"
)
func TestContextLimitForModel(t *testing.T) {
tests := []struct {
model string
expected int
}{
{"claude-3-opus-20240229", 200000},
{"claude-3-sonnet-20240229", 200000},
{"claude-3-haiku-20240307", 200000},
{"claude-sonnet-4-20250514", 200000},
{"CLAUDE-3-OPUS", 200000}, // case insensitive
{"some-sonnet-model", 200000}, // substring match
{"unknown-model", 200000}, // default
}
for _, tc := range tests {
t.Run(tc.model, func(t *testing.T) {
got := contextLimitForModel(tc.model)
if got != tc.expected {
t.Errorf("contextLimitForModel(%q) = %d, want %d", tc.model, got, tc.expected)
}
})
}
}
func TestClientReset(t *testing.T) {
// Create a client with dummy API key (won't make real calls)
c := NewClient("dummy-key")
// Manually set some state
c.messages = []Message{
{Role: "user", Content: "hello"},
{Role: "assistant", Content: "hi"},
}
c.lastTokens = 100
c.totalTokens = 500
// Reset should clear everything
c.Reset()
if len(c.messages) != 0 {
t.Errorf("Reset() should clear messages, got %d", len(c.messages))
}
if c.lastTokens != 0 {
t.Errorf("Reset() should clear lastTokens, got %d", c.lastTokens)
}
if c.totalTokens != 0 {
t.Errorf("Reset() should clear totalTokens, got %d", c.totalTokens)
}
}
func TestClientTotalTokens(t *testing.T) {
c := NewClient("dummy-key")
// Initially should be 0
if got := c.TotalTokens(); got != 0 {
t.Errorf("TotalTokens() = %d, want 0", got)
}
// Manually set for testing
c.totalTokens = 12345
if got := c.TotalTokens(); got != 12345 {
t.Errorf("TotalTokens() = %d, want 12345", got)
}
}
func TestClientContextLimit(t *testing.T) {
c := NewClient("dummy-key")
// Default model should have 200K limit
got := c.ContextLimit()
if got != 200000 {
t.Errorf("ContextLimit() = %d, want 200000", got)
}
// Change model and verify
c.SetModel("claude-3-haiku-20240307")
got = c.ContextLimit()
if got != 200000 {
t.Errorf("ContextLimit() for haiku = %d, want 200000", got)
}
}
func TestClientTemperature(t *testing.T) {
c := NewClient("dummy-key")
// Default temperature
if got := c.Temperature(); got != 0.7 {
t.Errorf("Temperature() = %f, want 0.7", got)
}
// Valid temperature
if err := c.SetTemperature(1.5); err != nil {
t.Errorf("SetTemperature(1.5) error: %v", err)
}
if got := c.Temperature(); got != 1.5 {
t.Errorf("Temperature() = %f, want 1.5", got)
}
// Invalid temperature - too low
if err := c.SetTemperature(-0.1); err == nil {
t.Error("SetTemperature(-0.1) should return error")
}
// Invalid temperature - too high
if err := c.SetTemperature(2.1); err == nil {
t.Error("SetTemperature(2.1) should return error")
}
}
func TestClientSystemPrompt(t *testing.T) {
c := NewClient("dummy-key")
// Initially empty
if got := c.SystemPrompt(); got != "" {
t.Errorf("SystemPrompt() = %q, want empty", got)
}
// Set and verify
c.SetSystemPrompt("You are a helpful assistant")
if got := c.SystemPrompt(); got != "You are a helpful assistant" {
t.Errorf("SystemPrompt() = %q, want 'You are a helpful assistant'", got)
}
}
func TestClientMessages(t *testing.T) {
c := NewClient("dummy-key")
// Initially empty
if msgs := c.Messages(); len(msgs) != 0 {
t.Errorf("Messages() should be empty, got %d", len(msgs))
}
// Add messages
c.messages = []Message{
{Role: "user", Content: "hello"},
{Role: "assistant", Content: "hi"},
}
msgs := c.Messages()
if len(msgs) != 2 {
t.Errorf("Messages() = %d messages, want 2", len(msgs))
}
// Verify it's a copy (modification doesn't affect original)
msgs[0].Content = "modified"
if c.messages[0].Content == "modified" {
t.Error("Messages() should return a copy, not the original")
}
}
func TestClientAddSystemMessage(t *testing.T) {
c := NewClient("dummy-key")
// Add system message
c.AddSystemMessage("Context info here")
msgs := c.Messages()
if len(msgs) != 1 {
t.Fatalf("Messages() = %d messages, want 1", len(msgs))
}
if msgs[0].Role != "system" {
t.Errorf("Message role = %q, want 'system'", msgs[0].Role)
}
if msgs[0].Content != "Context info here" {
t.Errorf("Message content = %q, want 'Context info here'", msgs[0].Content)
}
}
func TestClientModel(t *testing.T) {
c := NewClient("dummy-key")
// Default model
if got := c.Model(); got != "claude-sonnet-4-20250514" {
t.Errorf("Model() = %q, want 'claude-sonnet-4-20250514'", got)
}
// Change model
c.SetModel("claude-3-haiku-20240307")
if got := c.Model(); got != "claude-3-haiku-20240307" {
t.Errorf("Model() = %q, want 'claude-3-haiku-20240307'", got)
}
}
func TestClientMessagesJSON(t *testing.T) {
c := NewClient("dummy-key")
c.messages = []Message{
{Role: "user", Content: "hello"},
}
data, err := c.MessagesJSON()
if err != nil {
t.Fatalf("MessagesJSON() error: %v", err)
}
// Should contain the message content
json := string(data)
if !contains(json, "hello") || !contains(json, "user") {
t.Errorf("MessagesJSON() = %s, should contain 'hello' and 'user'", json)
}
}
func TestClientIsStreaming(t *testing.T) {
c := NewClient("dummy-key")
// Initially not streaming
if c.IsStreaming() {
t.Error("IsStreaming() should be false initially")
}
}
// Helper function
func contains(s, substr string) bool {
for i := 0; i <= len(s)-len(substr); i++ {
if s[i:i+len(substr)] == substr {
return true
}
}
return false
}