test(17-02): add failing tests for scheduler package
- Storage round-trip test for SaveScheduledJob/ListScheduledJobs - Subscriber round-trip test for Add/Remove/List/IsSubscribed - Scheduler Start loads enabled jobs from DB - Scheduler AddJob/RemoveJob persists and registers - Scheduler RunJob manual trigger with callback
This commit is contained in:
1
pkg/scheduler/jobs.go
Normal file
1
pkg/scheduler/jobs.go
Normal file
@@ -0,0 +1 @@
|
||||
package scheduler
|
||||
1
pkg/scheduler/scheduler.go
Normal file
1
pkg/scheduler/scheduler.go
Normal file
@@ -0,0 +1 @@
|
||||
package scheduler
|
||||
204
pkg/scheduler/scheduler_test.go
Normal file
204
pkg/scheduler/scheduler_test.go
Normal file
@@ -0,0 +1,204 @@
|
||||
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()
|
||||
}
|
||||
Reference in New Issue
Block a user