odidere/internal/db/queries_test.go

package db

import (
	"context"
	"errors"
	"log/slog"
	"testing"

	"github.com/mattn/go-sqlite3"
)

func TestQueries(t *testing.T) {
	var (
		ctx = t.Context()
		db  = createDB(t)
	)
	defer db.Close()

	if err := db.ApplySchema(ctx, slog.Default()); err != nil {
		t.Fatalf("ApplySchema: %v", err)
	}

	t.Run("User", func(t *testing.T) {
		testUser(t, ctx, db)
	})
	t.Run("Conversation", func(t *testing.T) {
		testConversation(t, ctx, db)
	})
	t.Run("Messaging", func(t *testing.T) {
		testMessaging(t, ctx, db)
	})
	t.Run("ConversationTimestamp", func(t *testing.T) {
		testConversationTimestamp(t, ctx, db)
	})
	t.Run("Executions", func(t *testing.T) {
		testExecutions(t, ctx, db)
	})
	t.Run("Delete", func(t *testing.T) {
		testDelete(t, ctx, db)
	})
}

func testUser(t *testing.T, ctx context.Context, db *SQLite) {
	var (
		qr = New(db.Reader)
		qw = New(db.Writer)
	)
	u, err := qw.InsertUser(ctx, InsertUserParams{
		Name:     "user",
		Password: "test",
		Role:     UserRoleOwner,
	})
	if err != nil {
		t.Fatalf("InsertUser: %v", err)
	}

	res, err := qr.SelectUserId(ctx, u.Id)
	if err != nil {
		t.Fatalf("SelectUserId: %v", err)
	}
	if res == nil {
		t.Fatal("SelectUserId returned nil")
	}
	if res.Name != u.Name {
		t.Errorf(
			"expected name %s, got %s",
			u.Name, res.Name,
		)
	}
}

func testConversation(t *testing.T, ctx context.Context, db *SQLite) {
	var (
		qr    = New(db.Reader)
		qw    = New(db.Writer)
		title = "test"
	)

	u, err := qw.InsertUser(ctx, InsertUserParams{
		Name:     "conversation",
		Password: "test",
		Role:     UserRoleOwner,
	})
	if err != nil {
		t.Fatalf("InsertUser: %v", err)
	}

	c, err := qw.InsertConversation(ctx, InsertConversationParams{
		UserId: u.Id,
		Title:  &title,
	})
	if err != nil {
		t.Fatalf("InsertConversation: %v", err)
	}

	res, err := qr.SelectConversationId(ctx, c.Id)
	if err != nil {
		t.Fatalf("SelectConversationId: %v", err)
	}
	if res == nil {
		t.Fatal("SelectConversationId returned nil")
	}
	if res.Title == nil || *res.Title != title {
		t.Errorf("expected title %s, got %v", title, res.Title)
	}

	owned, err := qr.SelectConversationIdUserId(ctx, c.Id, u.Id)
	if err != nil {
		t.Fatalf("SelectConversationIdUserId: %v", err)
	}
	if owned.Id != c.Id {
		t.Errorf("expected conversation %s, got %s", c.Id, owned.Id)
	}
	if _, err := qr.SelectConversationIdUserId(ctx, c.Id, "unknown"); err == nil {
		t.Fatal("expected error selecting conversation for wrong user")
	}
	if _, err := qr.SelectConversationIdUserId(ctx, "unknown", u.Id); err == nil {
		t.Fatal("expected error selecting unknown conversation")
	}

	// Test conversation list.
	conversations, err := qr.SelectConversationsUserId(ctx, u.Id)
	if err != nil {
		t.Fatalf("SelectConversationsUserId: %v", err)
	}
	if len(conversations) != 1 {
		t.Fatalf("expected 1 conversation, got %d", len(conversations))
	}
}

func testMessaging(t *testing.T, ctx context.Context, db *SQLite) {
	var qw = New(db.Writer)

	u, err := qw.InsertUser(ctx, InsertUserParams{
		Name:     "messaging",
		Password: "test",
		Role:     UserRoleOwner,
	})
	if err != nil {
		t.Fatalf("InsertUser: %v", err)
	}

	c, err := qw.InsertConversation(ctx, InsertConversationParams{
		UserId: u.Id,
	})
	if err != nil {
		t.Fatalf("InsertConversation: %v", err)
	}

	// Insert a system message.
	seq, err := qw.NextSequence(ctx, c.Id)
	if err != nil {
		t.Fatalf("NextSequence: %v", err)
	}
	_, err = qw.InsertMessage(ctx, InsertMessageParams{
		ConversationId: c.Id,
		Sequence:       seq,
		Role:           MessageRoleSystem,
		Content:        `{"role":"system","content":"You are helpful"}`,
	})
	if err != nil {
		t.Fatalf("InsertMessage: %v", err)
	}

	// Insert a user message.
	seq, err = qw.NextSequence(ctx, c.Id)
	if err != nil {
		t.Fatalf("NextSequence: %v", err)
	}
	_, err = qw.InsertMessage(ctx, InsertMessageParams{
		ConversationId: c.Id,
		Sequence:       seq,
		Role:           MessageRoleUser,
		Content:        `{"role":"user","content":"Hello"}`,
	})
	if err != nil {
		t.Fatalf("InsertMessage: %v", err)
	}

	// Insert an assistant message with tool calls.
	seq, err = qw.NextSequence(ctx, c.Id)
	if err != nil {
		t.Fatalf("NextSequence: %v", err)
	}
	_, err = qw.InsertMessage(ctx, InsertMessageParams{
		ConversationId: c.Id,
		Sequence:       seq,
		Role:           MessageRoleAssistant,
		Content:        `{"role":"assistant","content":null,"tool_calls":[{"id":"call_abc","type":"function","function":{"name":"get_weather","arguments":"{\"location\":\"NYC\"}"}}]}`,
	})
	if err != nil {
		t.Fatalf("InsertMessage: %v", err)
	}

	// Insert a tool result message.
	seq, err = qw.NextSequence(ctx, c.Id)
	if err != nil {
		t.Fatalf("NextSequence: %v", err)
	}
	toolMsg, err := qw.InsertMessage(ctx, InsertMessageParams{
		ConversationId: c.Id,
		Sequence:       seq,
		Role:           MessageRoleTool,
		Content:        `{"role":"tool","content":"{\"temperature\":72}","tool_call_id":"call_abc"}`,
	})
	if err != nil {
		t.Fatalf("InsertMessage: %v", err)
	}

	// Test SelectMessages.
	qr := New(db.Reader)
	messages, err := qr.SelectMessagesConversationId(ctx, c.Id)
	if err != nil {
		t.Fatalf("SelectMessagesConversationId: %v", err)
	}
	if len(messages) != 4 {
		t.Fatalf("expected 4 messages, got %d", len(messages))
	}
	if messages[0].Role != MessageRoleSystem {
		t.Errorf("expected first message role system, got %s", messages[0].Role)
	}
	if messages[3].Role != MessageRoleTool {
		t.Errorf("expected last message role tool, got %s", messages[3].Role)
	}

	// Test UpdateMessageContent.
	if err := qw.UpdateMessageContent(ctx, UpdateMessageContentParams{
		Id:      toolMsg.Id,
		Content: `{"role":"tool","content":"{\"temperature\":75}","tool_call_id":"call_abc"}`,
	}); err != nil {
		t.Fatalf("UpdateMessageContent: %v", err)
	}

	updated, err := qr.SelectMessagesConversationId(ctx, c.Id)
	if err != nil {
		t.Fatalf("SelectMessagesConversationId: %v", err)
	}
	if updated[3].Content != `{"role":"tool","content":"{\"temperature\":75}","tool_call_id":"call_abc"}` {
		t.Errorf("updated content mismatch: %s", updated[3].Content)
	}
}

func testConversationTimestamp(t *testing.T, ctx context.Context, database *SQLite) {
	var (
		qr = New(database.Reader)
		qw = New(database.Writer)
	)

	u, err := qw.InsertUser(ctx, InsertUserParams{
		Name:     "conversation-timestamp",
		Password: "test",
		Role:     UserRoleOwner,
	})
	if err != nil {
		t.Fatalf("InsertUser: %v", err)
	}
	c, err := qw.InsertConversation(ctx, InsertConversationParams{UserId: u.Id})
	if err != nil {
		t.Fatalf("InsertConversation: %v", err)
	}

	const oldTimestamp = "2000-01-01T00:00:00.000Z"
	if _, err := database.Writer.ExecContext(
		ctx,
		"UPDATE conversations SET updated_at = ? WHERE id = ?",
		oldTimestamp,
		c.Id,
	); err != nil {
		t.Fatalf("set conversation timestamp before insert: %v", err)
	}
	msg, err := qw.InsertMessage(ctx, InsertMessageParams{
		ConversationId: c.Id,
		Sequence:       0,
		Role:           MessageRoleUser,
		Content:        `{"role":"user","content":"Hello"}`,
	})
	if err != nil {
		t.Fatalf("InsertMessage: %v", err)
	}
	conv, err := qr.SelectConversationId(ctx, c.Id)
	if err != nil {
		t.Fatalf("SelectConversationId after message insert: %v", err)
	}
	if conv.UpdatedAt.Format("2006-01-02T15:04:05.000Z") == oldTimestamp {
		t.Error("message insert did not update conversation timestamp")
	}

	if _, err := database.Writer.ExecContext(
		ctx,
		"UPDATE conversations SET updated_at = ? WHERE id = ?",
		oldTimestamp,
		c.Id,
	); err != nil {
		t.Fatalf("set conversation timestamp before update: %v", err)
	}
	if err := qw.UpdateMessageContent(ctx, UpdateMessageContentParams{
		Id:      msg.Id,
		Content: `{"role":"user","content":"Updated"}`,
	}); err != nil {
		t.Fatalf("UpdateMessageContent: %v", err)
	}
	conv, err = qr.SelectConversationId(ctx, c.Id)
	if err != nil {
		t.Fatalf("SelectConversationId after message update: %v", err)
	}
	if conv.UpdatedAt.Format("2006-01-02T15:04:05.000Z") == oldTimestamp {
		t.Error("message update did not update conversation timestamp")
	}

	if _, err := database.Writer.ExecContext(
		ctx,
		"UPDATE conversations SET updated_at = ? WHERE id = ?",
		oldTimestamp,
		c.Id,
	); err != nil {
		t.Fatalf("set conversation timestamp before delete: %v", err)
	}
	if err := qw.DeleteMessageId(ctx, msg.Id); err != nil {
		t.Fatalf("DeleteMessageId: %v", err)
	}
	conv, err = qr.SelectConversationId(ctx, c.Id)
	if err != nil {
		t.Fatalf("SelectConversationId after message delete: %v", err)
	}
	if conv.UpdatedAt.Format("2006-01-02T15:04:05.000Z") == oldTimestamp {
		t.Error("message delete did not update conversation timestamp")
	}
}

func testDelete(t *testing.T, ctx context.Context, db *SQLite) {
	var (
		qr = New(db.Reader)
		qw = New(db.Writer)
	)

	u, err := qw.InsertUser(ctx, InsertUserParams{
		Name:     "delete",
		Password: "test",
		Role:     UserRoleOwner,
	})
	if err != nil {
		t.Fatalf("InsertUser: %v", err)
	}

	c, err := qw.InsertConversation(ctx, InsertConversationParams{
		UserId: u.Id,
	})
	if err != nil {
		t.Fatalf("InsertConversation: %v", err)
	}

	msg, err := qw.InsertMessage(ctx, InsertMessageParams{
		ConversationId: c.Id,
		Sequence:       1,
		Role:           MessageRoleUser,
		Content:        `{"role":"user","content":"Hello"}`,
	})
	if err != nil {
		t.Fatalf("InsertMessage: %v", err)
	}

	// Delete message.
	if err := qw.DeleteMessageId(ctx, msg.Id); err != nil {
		t.Fatalf("DeleteMessageId: %v", err)
	}

	messages, err := qr.SelectMessagesConversationId(ctx, c.Id)
	if err != nil {
		t.Fatalf("SelectMessagesConversationId: %v", err)
	}
	if len(messages) != 0 {
		t.Fatalf("expected 0 messages after delete, got %d", len(messages))
	}

	// Delete conversation.
	if err := qw.DeleteConversationId(ctx, c.Id); err != nil {
		t.Fatalf("DeleteConversationId: %v", err)
	}

	_, err = qr.SelectConversationId(ctx, c.Id)
	if err == nil {
		t.Fatal("expected error selecting deleted conversation")
	}
}

func testExecutions(t *testing.T, ctx context.Context, db *SQLite) {
	var (
		qr = New(db.Reader)
		qw = New(db.Writer)
	)

	u, err := qw.InsertUser(ctx, InsertUserParams{
		Name:     "approvals",
		Password: "test",
		Role:     UserRoleOwner,
	})
	if err != nil {
		t.Fatalf("InsertUser: %v", err)
	}

	c, err := qw.InsertConversation(ctx, InsertConversationParams{
		UserId: u.Id,
	})
	if err != nil {
		t.Fatalf("InsertConversation: %v", err)
	}

	msg, err := qw.InsertMessage(ctx, InsertMessageParams{
		ConversationId: c.Id,
		Sequence:       1,
		Role:           MessageRoleAssistant,
		Content:        `{"role":"assistant","content":null,"tool_calls":[{"id":"call_abc","type":"function","function":{"name":"host_run","arguments":"{\"command\":\"rm -rf /\"}"}}]}`,
	})
	if err != nil {
		t.Fatalf("InsertMessage: %v", err)
	}

	// Insert an execution.
	execution, err := qw.InsertExecution(ctx, InsertExecutionParams{
		MessageId:  msg.Id,
		Status:     ExecutionStatusPending,
		ToolCallId: "call_abc",
	})
	if err != nil {
		t.Fatalf("InsertExecution: %v", err)
	}
	if execution.Status != ExecutionStatusPending {
		t.Fatalf("expected pending, got %s", execution.Status)
	}
	if execution.ToolCallId != "call_abc" {
		t.Fatalf("expected tool_call_id call_abc, got %s", execution.ToolCallId)
	}

	// Select by ID.
	res, err := qr.SelectExecutionId(ctx, execution.Id)
	if err != nil {
		t.Fatalf("SelectExecutionId: %v", err)
	}
	if res == nil {
		t.Fatal("SelectExecutionId returned nil")
	}
	if res.MessageId != msg.Id {
		t.Errorf("expected message_id %s, got %s", msg.Id, res.MessageId)
	}

	// Select by message ID.
	executions, err := qr.SelectExecutionsMessageId(ctx, msg.Id)
	if err != nil {
		t.Fatalf("SelectExecutionsMessageId: %v", err)
	}
	if len(executions) != 1 {
		t.Fatalf("expected 1 execution, got %d", len(executions))
	}

	// Select pending.
	pending, err := qr.SelectPendingExecutions(ctx)
	if err != nil {
		t.Fatalf("SelectPendingExecutions: %v", err)
	}
	if len(pending) != 1 {
		t.Fatalf("expected 1 pending execution, got %d", len(pending))
	}

	// Transition to approved.
	if err := qw.UpdateExecutionStatus(ctx, UpdateExecutionStatusParams{
		Id:     execution.Id,
		Status: ExecutionStatusApproved,
	}); err != nil {
		t.Fatalf("UpdateExecutionStatus: %v", err)
	}

	res, err = qr.SelectExecutionId(ctx, execution.Id)
	if err != nil {
		t.Fatalf("SelectExecutionId after update: %v", err)
	}
	if res.Status != ExecutionStatusApproved {
		t.Fatalf("expected approved, got %s", res.Status)
	}

	// Pending should now be empty.
	pending, err = qr.SelectPendingExecutions(ctx)
	if err != nil {
		t.Fatalf("SelectPendingExecutions after update: %v", err)
	}
	if len(pending) != 0 {
		t.Fatalf("expected 0 pending, got %d", len(pending))
	}

	// Insert a second execution and deny it.
	execution2, err := qw.InsertExecution(ctx, InsertExecutionParams{
		MessageId:  msg.Id,
		Status:     ExecutionStatusPending,
		ToolCallId: "call_def",
	})
	if err != nil {
		t.Fatalf("InsertExecution second: %v", err)
	}

	if err := qw.UpdateExecutionStatus(ctx, UpdateExecutionStatusParams{
		Id:     execution2.Id,
		Status: ExecutionStatusDenied,
	}); err != nil {
		t.Fatalf("UpdateExecutionStatus denied: %v", err)
	}

	res2, err := qr.SelectExecutionId(ctx, execution2.Id)
	if err != nil {
		t.Fatalf("SelectExecutionId after deny: %v", err)
	}
	if res2.Status != ExecutionStatusDenied {
		t.Fatalf("expected denied, got %s", res2.Status)
	}

	// Transition to running, then completed.
	if err := qw.UpdateExecutionStatus(ctx, UpdateExecutionStatusParams{
		Id:     execution.Id,
		Status: ExecutionStatusRunning,
	}); err != nil {
		t.Fatalf("UpdateExecutionStatus running: %v", err)
	}
	if err := qw.UpdateExecutionStatus(ctx, UpdateExecutionStatusParams{
		Id:     execution.Id,
		Status: ExecutionStatusCompleted,
	}); err != nil {
		t.Fatalf("UpdateExecutionStatus completed: %v", err)
	}

	// Both executions for the message, ordered by created_at.
	all, err := qr.SelectExecutionsMessageId(ctx, msg.Id)
	if err != nil {
		t.Fatalf("SelectExecutionsMessageId final: %v", err)
	}
	if len(all) != 2 {
		t.Fatalf("expected 2 executions, got %d", len(all))
	}
	if all[0].Status != ExecutionStatusCompleted {
		t.Fatalf("expected first execution completed, got %s", all[0].Status)
	}
	if all[1].Status != ExecutionStatusDenied {
		t.Fatalf("expected second execution denied, got %s", all[1].Status)
	}

	// Parameter values and the schema default, on a second message.
	msg2, err := qw.InsertMessage(ctx, InsertMessageParams{
		ConversationId: c.Id,
		Sequence:       2,
		Role:           MessageRoleAssistant,
		Content:        `{"role":"assistant","content":"second"}`,
	})
	if err != nil {
		t.Fatalf("InsertMessage second: %v", err)
	}

	execRunning, err := qw.InsertExecution(ctx, InsertExecutionParams{
		MessageId:  msg2.Id,
		Status:     ExecutionStatusRunning,
		ToolCallId: "call_running",
	})
	if err != nil {
		t.Fatalf("InsertExecution running: %v", err)
	}
	if execRunning.Status != ExecutionStatusRunning {
		t.Fatalf("expected running, got %s", execRunning.Status)
	}

	// A status other than running is stored as passed.
	execPending, err := qw.InsertExecution(ctx, InsertExecutionParams{
		MessageId:  msg2.Id,
		Status:     ExecutionStatusPending,
		ToolCallId: "call_pending",
	})
	if err != nil {
		t.Fatalf("InsertExecution pending: %v", err)
	}
	if execPending.Status != ExecutionStatusPending {
		t.Fatalf("expected pending, got %s", execPending.Status)
	}

	// The zero value is rejected by the table's status check.
	_, err = qw.InsertExecution(ctx, InsertExecutionParams{
		MessageId:  msg2.Id,
		ToolCallId: "call_empty",
	})
	var sqliteErr sqlite3.Error
	if !errors.As(err, &sqliteErr) {
		t.Fatalf("InsertExecution with zero status: expected sqlite3.Error, got %v", err)
	}
	if sqliteErr.ExtendedCode != sqlite3.ErrConstraintCheck {
		t.Fatalf("InsertExecution with zero status: expected ErrConstraintCheck, got %v", err)
	}

	// A second row for the same message and tool call is rejected.
	_, err = qw.InsertExecution(ctx, InsertExecutionParams{
		MessageId:  msg2.Id,
		Status:     ExecutionStatusPending,
		ToolCallId: "call_running",
	})
	if !errors.As(err, &sqliteErr) {
		t.Fatalf("InsertExecution duplicate: expected sqlite3.Error, got %v", err)
	}
	if sqliteErr.ExtendedCode != sqlite3.ErrConstraintUnique {
		t.Fatalf("InsertExecution duplicate: expected ErrConstraintUnique, got %v", err)
	}

	// A row inserted without a status takes the schema default, running.
	if _, err := db.Writer.ExecContext(
		ctx,
		`INSERT INTO executions (id, message_id, tool_call_id) VALUES (?, ?, ?)`,
		"exec_default", msg2.Id, "call_default",
	); err != nil {
		t.Fatalf("insert execution without status: %v", err)
	}
	res, err = qr.SelectExecutionId(ctx, "exec_default")
	if err != nil {
		t.Fatalf("SelectExecutionId default: %v", err)
	}
	if res == nil {
		t.Fatal("SelectExecutionId default returned nil")
	}
	if res.Status != ExecutionStatusRunning {
		t.Fatalf("expected default running, got %s", res.Status)
	}

	// The rejected insert stored no row: the message's rows are exactly the
	// three successful inserts.
	all2, err := qr.SelectExecutionsMessageId(ctx, msg2.Id)
	if err != nil {
		t.Fatalf("SelectExecutionsMessageId second: %v", err)
	}
	seen := map[string]bool{}
	for _, e := range all2 {
		seen[e.ToolCallId] = true
	}
	for _, id := range []string{"call_default", "call_pending", "call_running"} {
		if !seen[id] {
			t.Fatalf("execution %s not found", id)
		}
	}
	if len(seen) != 3 {
		t.Fatalf("expected 3 executions, got %d", len(seen))
	}
}