package storage_test import ( "database/sql" "errors" "testing" "github.com/salvacybersec/keyhunter/pkg/storage" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) func sampleCustomDork() storage.CustomDork { return storage.CustomDork{ DorkID: "user-openai-envfile", Name: "User OpenAI .env", Source: "github", Category: "frontier", Query: "sk-proj- extension:env", Description: "user-added dork", Tags: []string{"openai", "env"}, } } func TestSaveCustomDork_RoundTrip(t *testing.T) { db, err := storage.Open(":memory:") require.NoError(t, err) defer db.Close() id, err := db.SaveCustomDork(sampleCustomDork()) require.NoError(t, err) assert.Greater(t, id, int64(0)) got, err := db.GetCustomDork(id) require.NoError(t, err) assert.Equal(t, "user-openai-envfile", got.DorkID) assert.Equal(t, "github", got.Source) assert.Equal(t, "frontier", got.Category) assert.Equal(t, "sk-proj- extension:env", got.Query) assert.Equal(t, "user-added dork", got.Description) assert.Equal(t, []string{"openai", "env"}, got.Tags) assert.False(t, got.CreatedAt.IsZero()) } func TestListCustomDorks_NewestFirst(t *testing.T) { db, err := storage.Open(":memory:") require.NoError(t, err) defer db.Close() first := sampleCustomDork() first.DorkID = "first" _, err = db.SaveCustomDork(first) require.NoError(t, err) second := sampleCustomDork() second.DorkID = "second" _, err = db.SaveCustomDork(second) require.NoError(t, err) list, err := db.ListCustomDorks() require.NoError(t, err) require.Len(t, list, 2) assert.Equal(t, "second", list[0].DorkID, "newest should be first") assert.Equal(t, "first", list[1].DorkID) } func TestGetCustomDork_NotFound(t *testing.T) { db, err := storage.Open(":memory:") require.NoError(t, err) defer db.Close() _, err = db.GetCustomDork(9999) require.Error(t, err) assert.True(t, errors.Is(err, sql.ErrNoRows)) } func TestGetCustomDorkByDorkID(t *testing.T) { db, err := storage.Open(":memory:") require.NoError(t, err) defer db.Close() _, err = db.SaveCustomDork(sampleCustomDork()) require.NoError(t, err) got, err := db.GetCustomDorkByDorkID("user-openai-envfile") require.NoError(t, err) assert.Equal(t, "user-openai-envfile", got.DorkID) _, err = db.GetCustomDorkByDorkID("missing") require.Error(t, err) assert.True(t, errors.Is(err, sql.ErrNoRows)) } func TestDeleteCustomDork(t *testing.T) { db, err := storage.Open(":memory:") require.NoError(t, err) defer db.Close() id, err := db.SaveCustomDork(sampleCustomDork()) require.NoError(t, err) n, err := db.DeleteCustomDork(id) require.NoError(t, err) assert.Equal(t, int64(1), n) _, err = db.GetCustomDork(id) assert.True(t, errors.Is(err, sql.ErrNoRows)) // Deleting again is a no-op (0 rows affected, no error). n, err = db.DeleteCustomDork(id) require.NoError(t, err) assert.Equal(t, int64(0), n) } func TestSaveCustomDork_UniqueDorkID(t *testing.T) { db, err := storage.Open(":memory:") require.NoError(t, err) defer db.Close() _, err = db.SaveCustomDork(sampleCustomDork()) require.NoError(t, err) // Duplicate dork_id must fail the UNIQUE constraint. _, err = db.SaveCustomDork(sampleCustomDork()) require.Error(t, err) } func TestSchemaMigration_Idempotent(t *testing.T) { db, err := storage.Open(":memory:") require.NoError(t, err) defer db.Close() // Re-running the CREATE TABLE IF NOT EXISTS on the same connection // must be a no-op. We verify by executing CREATE TABLE again manually. _, err = db.SQL().Exec(` CREATE TABLE IF NOT EXISTS custom_dorks ( id INTEGER PRIMARY KEY AUTOINCREMENT, dork_id TEXT NOT NULL UNIQUE, name TEXT NOT NULL, source TEXT NOT NULL, category TEXT NOT NULL, query TEXT NOT NULL, description TEXT, tags TEXT, created_at DATETIME DEFAULT CURRENT_TIMESTAMP ); `) assert.NoError(t, err) // And that pre-existing rows survive the idempotent re-exec. _, err = db.SaveCustomDork(sampleCustomDork()) require.NoError(t, err) list, err := db.ListCustomDorks() require.NoError(t, err) assert.Len(t, list, 1) }