package ai import ( "path/filepath" "testing" "github.com/vstockwell/wraith/internal/db" ) func setupConversationManager(t *testing.T) *ConversationManager { t.Helper() database, err := db.Open(filepath.Join(t.TempDir(), "test.db")) if err != nil { t.Fatalf("open db: %v", err) } if err := db.Migrate(database); err != nil { t.Fatalf("migrate: %v", err) } // Create the conversations table (002 migration) _, err = database.Exec(`CREATE TABLE IF NOT EXISTS conversations ( id TEXT PRIMARY KEY, title TEXT, model TEXT NOT NULL, messages TEXT NOT NULL DEFAULT '[]', tokens_in INTEGER DEFAULT 0, tokens_out INTEGER DEFAULT 0, created_at DATETIME DEFAULT CURRENT_TIMESTAMP, updated_at DATETIME DEFAULT CURRENT_TIMESTAMP)`) if err != nil { t.Fatalf("create conversations table: %v", err) } t.Cleanup(func() { database.Close() }) return NewConversationManager(database) } func TestCreateConversation(t *testing.T) { mgr := setupConversationManager(t) conv, err := mgr.Create("claude-sonnet-4-20250514") if err != nil { t.Fatalf("create: %v", err) } if conv.ID == "" { t.Error("expected non-empty ID") } if conv.Model != "claude-sonnet-4-20250514" { t.Errorf("expected model claude-sonnet-4-20250514, got %s", conv.Model) } if conv.Title != "New conversation" { t.Errorf("expected title 'New conversation', got %s", conv.Title) } if conv.TokensIn != 0 || conv.TokensOut != 0 { t.Errorf("expected zero tokens, got in=%d out=%d", conv.TokensIn, conv.TokensOut) } } func TestAddAndGetMessages(t *testing.T) { mgr := setupConversationManager(t) conv, err := mgr.Create("test-model") if err != nil { t.Fatalf("create: %v", err) } // Add a user message userMsg := Message{ Role: "user", Content: []ContentBlock{ {Type: "text", Text: "What is running on port 8080?"}, }, } if err := mgr.AddMessage(conv.ID, userMsg); err != nil { t.Fatalf("add user message: %v", err) } // Add an assistant message assistantMsg := Message{ Role: "assistant", Content: []ContentBlock{ {Type: "text", Text: "Let me check that for you."}, }, } if err := mgr.AddMessage(conv.ID, assistantMsg); err != nil { t.Fatalf("add assistant message: %v", err) } // Get messages messages, err := mgr.GetMessages(conv.ID) if err != nil { t.Fatalf("get messages: %v", err) } if len(messages) != 2 { t.Fatalf("expected 2 messages, got %d", len(messages)) } if messages[0].Role != "user" { t.Errorf("expected first message role 'user', got %s", messages[0].Role) } if messages[1].Role != "assistant" { t.Errorf("expected second message role 'assistant', got %s", messages[1].Role) } if messages[0].Content[0].Text != "What is running on port 8080?" { t.Errorf("unexpected message text: %s", messages[0].Content[0].Text) } } func TestListConversations(t *testing.T) { mgr := setupConversationManager(t) // Create multiple conversations _, err := mgr.Create("model-a") if err != nil { t.Fatalf("create 1: %v", err) } _, err = mgr.Create("model-b") if err != nil { t.Fatalf("create 2: %v", err) } list, err := mgr.List() if err != nil { t.Fatalf("list: %v", err) } if len(list) != 2 { t.Errorf("expected 2 conversations, got %d", len(list)) } } func TestDeleteConversation(t *testing.T) { mgr := setupConversationManager(t) conv, err := mgr.Create("test-model") if err != nil { t.Fatalf("create: %v", err) } if err := mgr.Delete(conv.ID); err != nil { t.Fatalf("delete: %v", err) } // Verify it's gone list, err := mgr.List() if err != nil { t.Fatalf("list: %v", err) } if len(list) != 0 { t.Errorf("expected 0 conversations after delete, got %d", len(list)) } // Delete non-existent should error if err := mgr.Delete("nonexistent"); err == nil { t.Error("expected error deleting non-existent conversation") } } func TestTokenUsageTracking(t *testing.T) { mgr := setupConversationManager(t) conv, err := mgr.Create("test-model") if err != nil { t.Fatalf("create: %v", err) } // Update token usage multiple times if err := mgr.UpdateTokenUsage(conv.ID, 100, 50); err != nil { t.Fatalf("update tokens 1: %v", err) } if err := mgr.UpdateTokenUsage(conv.ID, 200, 100); err != nil { t.Fatalf("update tokens 2: %v", err) } // Verify totals list, err := mgr.List() if err != nil { t.Fatalf("list: %v", err) } if len(list) != 1 { t.Fatalf("expected 1 conversation, got %d", len(list)) } if list[0].TokensIn != 300 { t.Errorf("expected 300 tokens in, got %d", list[0].TokensIn) } if list[0].TokensOut != 150 { t.Errorf("expected 150 tokens out, got %d", list[0].TokensOut) } } func TestGetMessagesNonExistent(t *testing.T) { mgr := setupConversationManager(t) _, err := mgr.GetMessages("nonexistent-id") if err == nil { t.Error("expected error for non-existent conversation") } } func TestAutoTitle(t *testing.T) { mgr := setupConversationManager(t) conv, err := mgr.Create("test-model") if err != nil { t.Fatalf("create: %v", err) } msg := Message{ Role: "user", Content: []ContentBlock{ {Type: "text", Text: "Check disk usage on server-01"}, }, } if err := mgr.AddMessage(conv.ID, msg); err != nil { t.Fatalf("add message: %v", err) } // Verify the title was auto-set list, err := mgr.List() if err != nil { t.Fatalf("list: %v", err) } if list[0].Title != "Check disk usage on server-01" { t.Errorf("expected auto-title, got %q", list[0].Title) } }