AetherForge/server/internal/scheduler/fleet_scheduler_test.go

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