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