keyhunter/pkg/scheduler/scheduler_test.go

package scheduler_test

import (
	"context"
	"sync"
	"testing"
	"time"

	"github.com/salvacybersec/keyhunter/pkg/scheduler"
	"github.com/salvacybersec/keyhunter/pkg/storage"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func openTestDB(t *testing.T) *storage.DB {
	t.Helper()
	db, err := storage.Open(":memory:")
	require.NoError(t, err)
	t.Cleanup(func() { db.Close() })
	return db
}

func TestStorageRoundTrip(t *testing.T) {
	db := openTestDB(t)

	id, err := db.SaveScheduledJob(storage.ScheduledJob{
		Name:           "nightly-scan",
		CronExpr:       "0 2 * * *",
		ScanCommand:    "/tmp/repos",
		NotifyTelegram: true,
		Enabled:        true,
	})
	require.NoError(t, err)
	assert.Greater(t, id, int64(0))

	jobs, err := db.ListScheduledJobs()
	require.NoError(t, err)
	require.Len(t, jobs, 1)
	assert.Equal(t, "nightly-scan", jobs[0].Name)
	assert.Equal(t, "0 2 * * *", jobs[0].CronExpr)
	assert.Equal(t, "/tmp/repos", jobs[0].ScanCommand)
	assert.True(t, jobs[0].NotifyTelegram)
	assert.True(t, jobs[0].Enabled)

	got, err := db.GetScheduledJob("nightly-scan")
	require.NoError(t, err)
	assert.Equal(t, "nightly-scan", got.Name)

	now := time.Now().UTC()
	next := now.Add(24 * time.Hour)
	err = db.UpdateJobLastRun("nightly-scan", now, &next)
	require.NoError(t, err)

	got2, err := db.GetScheduledJob("nightly-scan")
	require.NoError(t, err)
	require.NotNil(t, got2.LastRun)
	require.NotNil(t, got2.NextRun)

	n, err := db.DeleteScheduledJob("nightly-scan")
	require.NoError(t, err)
	assert.Equal(t, int64(1), n)

	jobs, err = db.ListScheduledJobs()
	require.NoError(t, err)
	assert.Empty(t, jobs)
}

func TestSubscriberRoundTrip(t *testing.T) {
	db := openTestDB(t)

	err := db.AddSubscriber(12345, "alice")
	require.NoError(t, err)

	subs, err := db.ListSubscribers()
	require.NoError(t, err)
	require.Len(t, subs, 1)
	assert.Equal(t, int64(12345), subs[0].ChatID)
	assert.Equal(t, "alice", subs[0].Username)

	ok, err := db.IsSubscribed(12345)
	require.NoError(t, err)
	assert.True(t, ok)

	ok, err = db.IsSubscribed(99999)
	require.NoError(t, err)
	assert.False(t, ok)

	n, err := db.RemoveSubscriber(12345)
	require.NoError(t, err)
	assert.Equal(t, int64(1), n)

	subs, err = db.ListSubscribers()
	require.NoError(t, err)
	assert.Empty(t, subs)
}

func TestSchedulerStartLoadsJobs(t *testing.T) {
	db := openTestDB(t)

	_, err := db.SaveScheduledJob(storage.ScheduledJob{
		Name: "job-a", CronExpr: "0 * * * *", ScanCommand: "/a", Enabled: true,
	})
	require.NoError(t, err)
	_, err = db.SaveScheduledJob(storage.ScheduledJob{
		Name: "job-b", CronExpr: "0 * * * *", ScanCommand: "/b", Enabled: true,
	})
	require.NoError(t, err)
	// Disabled job should not be registered
	_, err = db.SaveScheduledJob(storage.ScheduledJob{
		Name: "job-c", CronExpr: "0 * * * *", ScanCommand: "/c", Enabled: false,
	})
	require.NoError(t, err)

	s, err := scheduler.New(scheduler.Config{
		DB: db,
		ScanFunc: func(ctx context.Context, cmd string) (int, error) {
			return 0, nil
		},
	})
	require.NoError(t, err)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	err = s.Start(ctx)
	require.NoError(t, err)
	defer s.Stop()

	assert.Equal(t, 2, s.JobCount())
}

func TestSchedulerAddRemoveJob(t *testing.T) {
	db := openTestDB(t)

	s, err := scheduler.New(scheduler.Config{
		DB: db,
		ScanFunc: func(ctx context.Context, cmd string) (int, error) {
			return 0, nil
		},
	})
	require.NoError(t, err)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	err = s.Start(ctx)
	require.NoError(t, err)
	defer s.Stop()

	err = s.AddJob("test-job", "0 * * * *", "/test", true)
	require.NoError(t, err)
	assert.Equal(t, 1, s.JobCount())

	jobs, err := db.ListScheduledJobs()
	require.NoError(t, err)
	require.Len(t, jobs, 1)
	assert.Equal(t, "test-job", jobs[0].Name)

	err = s.RemoveJob("test-job")
	require.NoError(t, err)
	assert.Equal(t, 0, s.JobCount())

	jobs, err = db.ListScheduledJobs()
	require.NoError(t, err)
	assert.Empty(t, jobs)
}

func TestSchedulerRunJob(t *testing.T) {
	db := openTestDB(t)

	var mu sync.Mutex
	var scanCalled string
	var completeCalled bool

	s, err := scheduler.New(scheduler.Config{
		DB: db,
		ScanFunc: func(ctx context.Context, cmd string) (int, error) {
			mu.Lock()
			scanCalled = cmd
			mu.Unlock()
			return 5, nil
		},
		OnComplete: func(result scheduler.JobResult) {
			mu.Lock()
			completeCalled = true
			mu.Unlock()
		},
	})
	require.NoError(t, err)

	_, err = db.SaveScheduledJob(storage.ScheduledJob{
		Name: "manual-run", CronExpr: "0 * * * *", ScanCommand: "/manual", NotifyTelegram: true, Enabled: true,
	})
	require.NoError(t, err)

	result, err := s.RunJob(context.Background(), "manual-run")
	require.NoError(t, err)
	assert.Equal(t, "manual-run", result.JobName)
	assert.Equal(t, 5, result.FindingCount)

	mu.Lock()
	assert.Equal(t, "/manual", scanCalled)
	assert.True(t, completeCalled)
	mu.Unlock()
}