keyhunter/pkg/storage/custom_dorks_test.go

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