Files
claude-mnemonic/internal/db/sqlite/prompt_test.go
T
2025-12-19 02:17:02 +00:00

197 lines
5.6 KiB
Go

package sqlite
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func testPromptStore(t *testing.T) (*PromptStore, *Store, func()) {
t.Helper()
db, _, cleanup := testDB(t)
createAllTables(t, db)
store := newStoreFromDB(db)
promptStore := NewPromptStore(store)
return promptStore, store, cleanup
}
func TestPromptStore_SaveUserPromptWithMatches(t *testing.T) {
promptStore, store, cleanup := testPromptStore(t)
defer cleanup()
ctx := context.Background()
// Create a session first
seedSession(t, storeDB(store), "claude-1", "sdk-1", "test-project")
// Save a prompt
id, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", 1, "Help me fix this bug", 5)
require.NoError(t, err)
assert.Greater(t, id, int64(0))
// Verify it was saved
var count int
err = storeDB(store).QueryRow("SELECT COUNT(*) FROM user_prompts WHERE id = ?", id).Scan(&count)
require.NoError(t, err)
assert.Equal(t, 1, count)
}
func TestPromptStore_GetAllRecentUserPrompts(t *testing.T) {
promptStore, store, cleanup := testPromptStore(t)
defer cleanup()
ctx := context.Background()
// Create a session
seedSession(t, storeDB(store), "claude-1", "sdk-1", "test-project")
// Save multiple prompts
for i := 1; i <= 5; i++ {
_, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", i, "Prompt "+string(rune('A'+i-1)), i)
require.NoError(t, err)
time.Sleep(time.Millisecond) // Ensure different timestamps
}
// Get recent prompts
prompts, err := promptStore.GetAllRecentUserPrompts(ctx, 3)
require.NoError(t, err)
assert.Len(t, prompts, 3)
// Should be in descending order (most recent first)
assert.Equal(t, 5, prompts[0].PromptNumber)
}
func TestPromptStore_GetRecentUserPromptsByProject(t *testing.T) {
promptStore, store, cleanup := testPromptStore(t)
defer cleanup()
ctx := context.Background()
// Create sessions for different projects
seedSession(t, storeDB(store), "claude-1", "sdk-1", "project-a")
seedSession(t, storeDB(store), "claude-2", "sdk-2", "project-b")
// Save prompts for both projects
for i := 1; i <= 3; i++ {
_, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", i, "Project A prompt", 0)
require.NoError(t, err)
}
for i := 1; i <= 2; i++ {
_, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-2", i, "Project B prompt", 0)
require.NoError(t, err)
}
// Get prompts for project-a
prompts, err := promptStore.GetRecentUserPromptsByProject(ctx, "project-a", 10)
require.NoError(t, err)
assert.Len(t, prompts, 3)
// Get prompts for project-b
prompts, err = promptStore.GetRecentUserPromptsByProject(ctx, "project-b", 10)
require.NoError(t, err)
assert.Len(t, prompts, 2)
}
func TestPromptStore_CleanupOldPrompts(t *testing.T) {
promptStore, store, cleanup := testPromptStore(t)
defer cleanup()
ctx := context.Background()
// Create a session
seedSession(t, storeDB(store), "claude-1", "sdk-1", "test-project")
// Save more prompts than the limit
// Note: MaxPromptsGlobal is 500, but we'll test with a smaller number
// by directly calling CleanupOldPrompts
for i := 1; i <= 10; i++ {
_, err := storeDB(store).Exec(`
INSERT INTO user_prompts (claude_session_id, prompt_number, prompt_text, created_at, created_at_epoch)
VALUES (?, ?, ?, datetime('now'), ?)
`, "claude-1", i, "Prompt "+string(rune('A'+i-1)), time.Now().UnixMilli()+int64(i))
require.NoError(t, err)
}
// Verify we have 10 prompts
var count int
err := storeDB(store).QueryRow("SELECT COUNT(*) FROM user_prompts").Scan(&count)
require.NoError(t, err)
assert.Equal(t, 10, count)
// Cleanup should return empty since we're under the limit
deletedIDs, err := promptStore.CleanupOldPrompts(ctx)
require.NoError(t, err)
assert.Empty(t, deletedIDs)
}
func TestPromptStore_SetCleanupFunc(t *testing.T) {
promptStore, store, cleanup := testPromptStore(t)
defer cleanup()
ctx := context.Background()
// Track cleanup calls
var cleanupCalledWith []int64
promptStore.SetCleanupFunc(func(ctx context.Context, deletedIDs []int64) {
cleanupCalledWith = deletedIDs
})
// Create a session
seedSession(t, storeDB(store), "claude-1", "sdk-1", "test-project")
// Save a prompt (should trigger cleanup, but won't delete anything under limit)
_, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", 1, "Test prompt", 0)
require.NoError(t, err)
// Cleanup func should not have been called since nothing was deleted
assert.Empty(t, cleanupCalledWith)
}
func TestPromptStore_GetPromptsByIDs(t *testing.T) {
promptStore, store, cleanup := testPromptStore(t)
defer cleanup()
ctx := context.Background()
// Create a session
seedSession(t, storeDB(store), "claude-1", "sdk-1", "test-project")
// Save some prompts and collect their IDs
var ids []int64
for i := 1; i <= 5; i++ {
id, err := promptStore.SaveUserPromptWithMatches(ctx, "claude-1", i, "Prompt "+string(rune('A'+i-1)), 0)
require.NoError(t, err)
ids = append(ids, id)
time.Sleep(time.Millisecond)
}
// Get specific prompts by ID
prompts, err := promptStore.GetPromptsByIDs(ctx, ids[:3], "date_desc", 10)
require.NoError(t, err)
assert.Len(t, prompts, 3)
// Test with ascending order
prompts, err = promptStore.GetPromptsByIDs(ctx, ids, "date_asc", 2)
require.NoError(t, err)
assert.Len(t, prompts, 2)
assert.Equal(t, 1, prompts[0].PromptNumber)
}
func TestPromptStore_GetPromptsByIDs_EmptyInput(t *testing.T) {
promptStore, _, cleanup := testPromptStore(t)
defer cleanup()
ctx := context.Background()
// Empty IDs should return nil
prompts, err := promptStore.GetPromptsByIDs(ctx, []int64{}, "date_desc", 10)
require.NoError(t, err)
assert.Nil(t, prompts)
}