package scheduler import ( "strings" "testing" "time" "crypto-miner-server/internal/db" "crypto-miner-server/internal/models" ) type mockSender struct { commands []string agents []string err error } func (m *mockSender) SendAgentCommand(agentID, action string, args map[string]interface{}) error { cmd := agentID + ":" + action if c, ok := args["command"].(string); ok && c != "" { cmd += ":" + c } m.commands = append(m.commands, cmd) return m.err } func (m *mockSender) ConnectedAgentIDs() []string { return m.agents } func openTestDB(t *testing.T) *db.Database { t.Helper() d, err := db.New(t.TempDir()) if err != nil { t.Fatal(err) } t.Cleanup(func() { d.Close() }) return d } func upsertTask(t *testing.T, d *db.Database, task *models.FleetTask) *models.FleetTask { t.Helper() if err := d.UpsertFleetTask(task); err != nil { t.Fatal(err) } list, err := d.ListFleetTasks() if err != nil || len(list) == 0 { t.Fatalf("list tasks: %v", err) } return list[0] } func TestRunConnectTasksDispatchesMatchingTrigger(t *testing.T) { d := openTestDB(t) send := &mockSender{agents: []string{"agent-a"}} s := New(d, send) task := upsertTask(t, d, &models.FleetTask{ Name: "sysinfo on connect", Enabled: true, Trigger: "on_connect", Action: "sysinfo", }) s.RunConnectTasks("agent-a", "on_connect") if len(send.commands) != 1 { t.Fatalf("commands = %v, want 1 dispatch", send.commands) } want := "agent-a:sysinfo" if send.commands[0] != want { t.Fatalf("got %q, want %q", send.commands[0], want) } last, ok := d.LastFleetTaskRun("agent-a", task.ID) if !ok || time.Since(last) > time.Minute { t.Fatalf("expected recent fleet task run record, ok=%v last=%v", ok, last) } } func TestRunConnectTasksSkipsDisabledAndWrongTrigger(t *testing.T) { d := openTestDB(t) send := &mockSender{} s := New(d, send) upsertTask(t, d, &models.FleetTask{ Name: "disabled", Enabled: false, Trigger: "on_connect", Action: "sysinfo", }) upsertTask(t, d, &models.FleetTask{ Name: "reconnect only", Enabled: true, Trigger: "on_reconnect", Action: "sysinfo", }) s.RunConnectTasks("agent-a", "on_connect") if len(send.commands) != 0 { t.Fatalf("expected no dispatch, got %v", send.commands) } } func TestRunConnectTasksIncludesCommandArg(t *testing.T) { d := openTestDB(t) send := &mockSender{} s := New(d, send) upsertTask(t, d, &models.FleetTask{ Name: "shell job", Enabled: true, Trigger: "on_connect", Action: "run_shell", Command: "whoami", }) s.RunConnectTasks("agent-b", "on_connect") if len(send.commands) != 1 || !strings.Contains(send.commands[0], ":whoami") { t.Fatalf("expected command arg in dispatch, got %v", send.commands) } } func TestTickIntervalRespectsLastRun(t *testing.T) { d := openTestDB(t) send := &mockSender{agents: []string{"agent-a"}} s := New(d, send) task := upsertTask(t, d, &models.FleetTask{ Name: "hourly sysinfo", Enabled: true, Trigger: "interval_hours", IntervalHours: 1, Action: "sysinfo", }) if err := d.RecordFleetTaskRun("agent-a", task.ID); err != nil { t.Fatal(err) } s.tickInterval() if len(send.commands) != 0 { t.Fatalf("expected skip within interval, got %v", send.commands) } s.tickInterval() } func TestTickIntervalDispatchesWhenDue(t *testing.T) { d := openTestDB(t) send := &mockSender{agents: []string{"agent-a", "agent-b"}} s := New(d, send) upsertTask(t, d, &models.FleetTask{ Name: "quick poll", Enabled: true, Trigger: "interval_hours", IntervalHours: 0.0001, Action: "heartbeat", }) s.tickInterval() if len(send.commands) != 2 { t.Fatalf("expected dispatch to both agents, got %v", send.commands) } } func TestTickCronDedupesSameDaySlot(t *testing.T) { d := openTestDB(t) send := &mockSender{agents: []string{"agent-a"}} s := New(d, send) slot := time.Now().Format("15:04") upsertTask(t, d, &models.FleetTask{ Name: "daily sysinfo", Enabled: true, Trigger: "cron", CronTime: slot, Action: "sysinfo", }) s.tickCron() first := len(send.commands) s.tickCron() second := len(send.commands) if first != 1 { t.Fatalf("first cron tick: got %d commands, want 1", first) } if second != first { t.Fatalf("cron dedupe failed: first=%d second=%d", first, second) } }