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