odidere/internal/db/queries_test.go

package db

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

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("Approvals", func(t *testing.T) {
		testApprovals(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 testApprovals(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 approval.
	approval, err := qw.InsertApproval(ctx, InsertApprovalParams{
		MessageId:  msg.Id,
		ToolCallId: "call_abc",
	})
	if err != nil {
		t.Fatalf("InsertApproval: %v", err)
	}
	if approval.Status != ApprovalStatusPending {
		t.Fatalf("expected pending, got %s", approval.Status)
	}
	if approval.ToolCallId != "call_abc" {
		t.Fatalf("expected tool_call_id call_abc, got %s", approval.ToolCallId)
	}

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

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

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

	// Transition to approved.
	if err := qw.UpdateApprovalStatus(ctx, UpdateApprovalStatusParams{
		Id:     approval.Id,
		Status: ApprovalStatusApproved,
	}); err != nil {
		t.Fatalf("UpdateApprovalStatus: %v", err)
	}

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

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

	// Insert a second approval and deny it.
	approval2, err := qw.InsertApproval(ctx, InsertApprovalParams{
		MessageId:  msg.Id,
		ToolCallId: "call_def",
	})
	if err != nil {
		t.Fatalf("InsertApproval second: %v", err)
	}

	if err := qw.UpdateApprovalStatus(ctx, UpdateApprovalStatusParams{
		Id:     approval2.Id,
		Status: ApprovalStatusDenied,
	}); err != nil {
		t.Fatalf("UpdateApprovalStatus denied: %v", err)
	}

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

	// Both approvals for the message, ordered by created_at.
	all, err := qr.SelectApprovalsMessageId(ctx, msg.Id)
	if err != nil {
		t.Fatalf("SelectApprovalsMessageId final: %v", err)
	}
	if len(all) != 2 {
		t.Fatalf("expected 2 approvals, got %d", len(all))
	}
	if all[0].Status != ApprovalStatusApproved {
		t.Fatalf("expected first approval approved, got %s", all[0].Status)
	}
	if all[1].Status != ApprovalStatusDenied {
		t.Fatalf("expected second approval denied, got %s", all[1].Status)
	}
}