523 lines
13 KiB
Go
523 lines
13 KiB
Go
package db
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
)
|
|
|
|
func TestQueries(t *testing.T) {
|
|
var (
|
|
ctx = t.Context()
|
|
db = createDB(t)
|
|
)
|
|
defer db.Close()
|
|
|
|
if err := db.ApplySchema(ctx, nil); 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)
|
|
}
|
|
}
|