odidere/internal/service/conversations_test.go

package service

import (
	"database/sql"
	"encoding/json/v2"
	"errors"
	"log/slog"
	"path/filepath"
	"sync"
	"testing"

	"code.chimeric.al/chimerical/odidere/internal/conversation"
	"code.chimeric.al/chimerical/odidere/internal/db"
	"code.chimeric.al/chimerical/odidere/internal/llm"
	"code.chimeric.al/chimerical/odidere/internal/tool"
)

// newTestConversationService returns a service with a fresh database.
func newTestConversationService(t *testing.T) *Service {
	t.Helper()
	var (
		ctx = t.Context()
		log = slog.New(slog.DiscardHandler)
	)
	database, err := db.Connect(ctx, db.Config{
		Path: filepath.Join(t.TempDir(), "conversations.db"),
	})
	if err != nil {
		t.Fatalf("Connect: %v", err)
	}
	t.Cleanup(func() { database.Close() })
	if err := database.ApplySchema(ctx, log); err != nil {
		t.Fatalf("ApplySchema: %v", err)
	}
	tools, err := tool.NewRegistry(nil)
	if err != nil {
		t.Fatalf("NewRegistry: %v", err)
	}
	return &Service{
		db:            database,
		dbr:           db.New(database.Reader),
		dbw:           db.New(database.Writer),
		log:           log,
		tools:         tools,
		conversations: make(map[string]*conversation.Conversation),
	}
}

func TestNewConversation(t *testing.T) {
	title := "Test conversation"
	tests := []struct {
		name   string
		title  *string
		system string
	}{
		{name: "with system", title: &title, system: "Use short answers."},
		{name: "empty system"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			svc := newTestConversationService(t)
			ctx := t.Context()
			saved, c, err := svc.newConversation(
				ctx, "default", tt.title, tt.system,
			)
			if err != nil {
				t.Fatalf("newConversation: %v", err)
			}
			if saved == nil || c == nil {
				t.Fatal("newConversation returned a nil row or handle")
			}
			stored, err := svc.dbr.SelectConversationId(ctx, saved.Id)
			if err != nil {
				t.Fatalf("SelectConversationId: %v", err)
			}
			if stored.UserId != "default" {
				t.Fatalf("UserId = %q, want default", stored.UserId)
			}
			if tt.title == nil {
				if stored.Title != nil {
					t.Fatalf("Title = %v, want nil", stored.Title)
				}
			} else if stored.Title == nil || *stored.Title != *tt.title {
				t.Fatalf("Title = %v, want %q", stored.Title, *tt.title)
			}
			rows, err := svc.dbr.SelectMessagesConversationId(ctx, saved.Id)
			if err != nil {
				t.Fatalf("SelectMessagesConversationId: %v", err)
			}
			if len(rows) != 1 {
				t.Fatalf("got %d messages, want 1", len(rows))
			}
			if rows[0].Role != db.MessageRoleSystem || rows[0].Sequence != 0 {
				t.Fatalf("unexpected system row: %+v", rows[0])
			}
			var msg llm.Message
			if err := json.Unmarshal(rows[0].Content, &msg); err != nil {
				t.Fatalf("Unmarshal: %v", err)
			}
			if err := msg.Validate(); err != nil {
				t.Fatalf("Validate: %v", err)
			}
			if msg.Role != llm.RoleSystem || msg.Text() != tt.system {
				t.Fatalf("unexpected system message: %+v", msg)
			}
			if svc.conversations[saved.Id] != c {
				t.Fatal("created handle was not registered")
			}
			got, err := svc.getConversation(ctx, saved.Id)
			if err != nil {
				t.Fatalf("getConversation: %v", err)
			}
			if got != c {
				t.Fatal("getConversation returned a different handle")
			}
		})
	}
}

func TestGetConversation(t *testing.T) {
	svc := newTestConversationService(t)
	ctx := t.Context()
	saved, err := svc.dbw.InsertConversation(
		ctx, db.InsertConversationParams{UserId: "default"},
	)
	if err != nil {
		t.Fatalf("InsertConversation: %v", err)
	}

	// Concurrent retrievals must share one handle for the saved row.
	var (
		wg      sync.WaitGroup
		handles [8]*conversation.Conversation
		errs    [8]error
	)
	for i := range handles {
		wg.Go(func() {
			handles[i], errs[i] = svc.getConversation(ctx, saved.Id)
		})
	}
	wg.Wait()
	for i, c := range handles {
		if errs[i] != nil {
			t.Fatalf("getConversation: %v", errs[i])
		}
		if c == nil || c != handles[0] || c != svc.conversations[saved.Id] {
			t.Fatal("getConversation did not reuse the registered handle")
		}
	}

	// A cached handle must not hide deletion of its saved row.
	if err := svc.dbw.DeleteConversationId(ctx, saved.Id); err != nil {
		t.Fatalf("DeleteConversationId: %v", err)
	}
	for _, id := range []string{saved.Id, "missing"} {
		c, err := svc.getConversation(ctx, id)
		if !errors.Is(err, sql.ErrNoRows) || c != nil {
			t.Fatalf("getConversation(%q) = %v, %v; want nil, sql.ErrNoRows",
				id, c, err)
		}
	}
	if _, ok := svc.conversations["missing"]; ok {
		t.Fatal("getConversation registered a missing conversation")
	}
}

func TestNewConversationRollback(t *testing.T) {
	for _, name := range []string{"missing owner", "missing logger", "missing tools"} {
		t.Run(name, func(t *testing.T) {
			svc := newTestConversationService(t)
			ctx := t.Context()
			userId := "default"
			switch name {
			case "missing owner":
				userId = "missing"
			case "missing logger":
				svc.log = nil
			case "missing tools":
				svc.tools = nil
			}
			saved, c, err := svc.newConversation(ctx, userId, nil, "System")
			if err == nil || saved != nil || c != nil {
				t.Fatalf("newConversation = %v, %v, %v; want nil, nil, error",
					saved, c, err)
			}
			var conversations, messages int
			if err := svc.db.Reader.QueryRowContext(ctx, `
SELECT
    (SELECT count(*) FROM conversations),
    (SELECT count(*) FROM messages)
`).Scan(&conversations, &messages); err != nil {
				t.Fatalf("count rows: %v", err)
			}
			if conversations != 0 || messages != 0 || len(svc.conversations) != 0 {
				t.Fatalf("got %d conversations, %d messages, %d handles; want none",
					conversations, messages, len(svc.conversations))
			}
		})
	}
}