wraith/internal/ai/conversation_test.go

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)
	}
}