fetch_ml/tests/unit/scheduler/failure_scenarios_test.go

package scheduler_test

import (
	"encoding/json"
	"testing"
	"time"

	"github.com/jfraeys/fetch_ml/internal/scheduler"
	fixtures "github.com/jfraeys/fetch_ml/tests/fixtures"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func mustMarshal(v any) []byte {
	b, _ := json.Marshal(v)
	return b
}

// TestWorkerDeath simulates a worker dying mid-job
func TestWorkerDeath_MidJob(t *testing.T) {
	// Use fixture for hub setup
	fixture := fixtures.NewSchedulerTestFixture(t, fixtures.DefaultHubConfig())
	defer fixture.Cleanup()

	// Create mock worker
	worker := fixture.CreateWorker("worker-death-test", scheduler.WorkerCapabilities{GPUCount: 0})

	// Send heartbeat
	worker.SendHeartbeat(scheduler.SlotStatus{
		BatchTotal: 3,
		BatchInUse: 1,
	})

	// Simulate worker death
	worker.Close()

	// Verify disconnect
	require.True(t, worker.WaitForDisconnect(2*time.Second), "worker should disconnect")
	t.Log("Worker death simulation completed")
}

// TestSchedulerRestartRecovery simulates scheduler restart
func TestSchedulerRestart_Recovery(t *testing.T) {
	dir := t.TempDir()

	// Create initial state store
	ss1, err := scheduler.NewStateStore(dir + "/state.json")
	require.NoError(t, err)

	// Record some events
	events := []scheduler.StateEvent{
		{Type: scheduler.EventJobEnqueued, TaskID: "task-1", Timestamp: time.Now()},
		{Type: scheduler.EventJobAssigned, TaskID: "task-1", WorkerID: "worker-1", Timestamp: time.Now()},
	}
	for _, e := range events {
		require.NoError(t, ss1.Append(e))
	}

	// Simulate restart by creating new state store
	ss2, err := scheduler.NewStateStore(dir + "/state.json")
	require.NoError(t, err)

	// Replay should recover state
	replayed, err := ss2.Replay()
	require.NoError(t, err)
	assert.Len(t, replayed, 2)
}

// TestSplitBrain_Case1: Worker reconnects with unknown task
func TestSplitBrain_UnknownTask(t *testing.T) {
	fixture := fixtures.NewSchedulerTestFixture(t, fixtures.DefaultHubConfig())
	defer fixture.Cleanup()

	// Create mock worker
	worker := fixture.CreateWorker("worker-split-1", scheduler.WorkerCapabilities{GPUCount: 0})

	// Simulate reconnect with unknown task
	worker.Send(scheduler.Message{
		Type: scheduler.MsgRegister,
		Payload: mustMarshal(scheduler.WorkerRegistration{
			ID: "worker-split-1",
			ActiveTasks: []scheduler.ActiveTaskReport{
				{TaskID: "unknown-task", State: "running"},
			},
		}),
	})

	// Should receive cancel for unknown task
	msg := worker.RecvTimeout(2 * time.Second)
	if msg.Type == scheduler.MsgJobCancel {
		t.Log("Received expected cancel for unknown task")
	} else {
		t.Logf("Received message type: %s (may need to check split-brain handling)", msg.Type)
	}
}

// TestSplitBrain_Case2: Worker reconnects with orphaned task
func TestSplitBrain_OrphanedTask(t *testing.T) {
	fixture := fixtures.NewSchedulerTestFixture(t, fixtures.DefaultHubConfig())
	defer fixture.Cleanup()

	// Create mock worker
	worker := fixture.CreateWorker("worker-split-2", scheduler.WorkerCapabilities{GPUCount: 0})

	// Simulate reconnect with orphaned task
	worker.Send(scheduler.Message{
		Type: scheduler.MsgRegister,
		Payload: mustMarshal(scheduler.WorkerRegistration{
			ID: "worker-split-2",
			ActiveTasks: []scheduler.ActiveTaskReport{
				{TaskID: "orphaned-task", State: "running"},
			},
		}),
	})

	msg := worker.RecvTimeout(2 * time.Second)
	t.Logf("Received message type: %s", msg.Type)
}

// TestSplitBrain_Case3: Worker reconnects with re-queued task
func TestSplitBrain_RequeuedTask(t *testing.T) {
	fixture := fixtures.NewSchedulerTestFixture(t, fixtures.DefaultHubConfig())
	defer fixture.Cleanup()

	// Create mock worker
	worker := fixture.CreateWorker("worker-split-3", scheduler.WorkerCapabilities{GPUCount: 0})

	// Simulate reconnect with re-queued task
	worker.Send(scheduler.Message{
		Type: scheduler.MsgRegister,
		Payload: mustMarshal(scheduler.WorkerRegistration{
			ID: "worker-split-3",
			ActiveTasks: []scheduler.ActiveTaskReport{
				{TaskID: "requeued-task", State: "queued"},
			},
		}),
	})

	msg := worker.RecvTimeout(2 * time.Second)
	t.Logf("Received message type: %s", msg.Type)
}

// TestAcceptanceTimeout: Job assigned but never accepted
func TestAcceptanceTimeout(t *testing.T) {
	cfg := fixtures.DefaultHubConfig()
	cfg.AcceptanceTimeoutSecs = 1
	fixture := fixtures.NewSchedulerTestFixture(t, cfg)
	defer fixture.Cleanup()

	// Create mock worker
	worker := fixture.CreateWorker("worker-timeout", scheduler.WorkerCapabilities{GPUCount: 0})

	// Signal ready but don't accept any job
	worker.SignalReady(scheduler.SlotStatus{BatchTotal: 3, BatchInUse: 0}, "polling")

	// Wait for potential assignment (but don't accept it)
	msg, ok := worker.RecvNonBlock()
	if ok && msg.Type == scheduler.MsgJobAssign {
		t.Log("Received job assignment, not accepting to test timeout")
	}

	// Wait for acceptance timeout
	time.Sleep(2 * time.Second)

	t.Log("Acceptance timeout test completed")
}

// TestGangTimeout: Multi-node job timeout during gang commit
func TestGangTimeout(t *testing.T) {
	cfg := fixtures.DefaultHubConfig()
	cfg.GangAllocTimeoutSecs = 1
	fixture := fixtures.NewSchedulerTestFixture(t, cfg)
	defer fixture.Cleanup()

	// Create mock worker
	worker := fixture.CreateWorker("worker-gang", scheduler.WorkerCapabilities{GPUCount: 0})

	// Signal ready
	worker.SignalReady(scheduler.SlotStatus{BatchTotal: 3, BatchInUse: 0}, "polling")

	// Wait for potential assignment
	_, _ = worker.RecvNonBlock()

	// Wait for gang timeout
	time.Sleep(2 * time.Second)

	t.Log("Gang timeout test completed")
}