poimen-workflows/internal/recovery/checkpoint_test.go

package recovery

import (
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func TestCheckpointManager(t *testing.T) {
	tmpDir := t.TempDir()

	cm := NewCheckpointManager(tmpDir, 100*time.Millisecond)
	err := cm.Start("wf-1")
	assert.NoError(t, err)
	defer cm.Stop()

	// Update stage
	err = cm.UpdateStage("clone")
	assert.NoError(t, err)

	// Add completed task
	err = cm.AddCompletedTask("task-1")
	assert.NoError(t, err)

	// Add pending tasks
	err = cm.SetPendingTasks([]string{"task-2", "task-3"})
	assert.NoError(t, err)

	// Get current checkpoint
	cp := cm.GetCurrent()
	assert.NotNil(t, cp)
	assert.Equal(t, "clone", cp.Stage)
	assert.Equal(t, 1, len(cp.CompletedTasks))
	assert.Equal(t, 2, len(cp.PendingTasks))
}

func TestCheckpointPersistence(t *testing.T) {
	tmpDir := t.TempDir()

	// Create and save checkpoint
	cm1 := NewCheckpointManager(tmpDir, 100*time.Millisecond)
	err := cm1.Start("wf-1")
	assert.NoError(t, err)

	cm1.UpdateStage("plan")
	cm1.AddCompletedTask("task-1")
	cm1.SetPendingTasks([]string{"task-2"})

	time.Sleep(150 * time.Millisecond) // Wait for periodic save
	cm1.Stop()

	// Load from disk
	cm2 := NewCheckpointManager(tmpDir, 100*time.Millisecond)
	cp, err := cm2.GetLatest("wf-1")
	assert.NoError(t, err)
	assert.NotNil(t, cp)
	assert.Equal(t, "plan", cp.Stage)
	assert.Equal(t, 1, len(cp.CompletedTasks))
}

func TestCheckpointHasCheckpoint(t *testing.T) {
	tmpDir := t.TempDir()

	cm := NewCheckpointManager(tmpDir, 100*time.Millisecond)
	err := cm.Start("wf-1")
	assert.NoError(t, err)
	defer cm.Stop()

	time.Sleep(150 * time.Millisecond)

	has, err := cm.HasCheckpoint("wf-1")
	assert.NoError(t, err)
	assert.True(t, has)

	has, err = cm.HasCheckpoint("wf-nonexistent")
	assert.NoError(t, err)
	assert.False(t, has)
}

func TestCheckpointCleanup(t *testing.T) {
	tmpDir := t.TempDir()

	cm := NewCheckpointManager(tmpDir, 100*time.Millisecond)
	err := cm.Start("wf-1")
	assert.NoError(t, err)

	time.Sleep(150 * time.Millisecond)
	cm.Stop()

	// Verify checkpoint exists
	has, err := cm.HasCheckpoint("wf-1")
	assert.NoError(t, err)
	assert.True(t, has)

	// Cleanup
	err = cm.CleanupCheckpoint("wf-1")
	assert.NoError(t, err)

	// Verify it's gone
	has, err = cm.HasCheckpoint("wf-1")
	assert.NoError(t, err)
	assert.False(t, has)
}

func TestCheckpointMetadata(t *testing.T) {
	tmpDir := t.TempDir()

	cm := NewCheckpointManager(tmpDir, 100*time.Millisecond)
	err := cm.Start("wf-1")
	assert.NoError(t, err)
	defer cm.Stop()

	// Add metadata
	cp := cm.GetCurrent()
	cp.Metadata["key"] = "value"
	cm.Update(cp)

	// Retrieve and verify
	retrieved := cm.GetCurrent()
	assert.Equal(t, "value", retrieved.Metadata["key"])
}

func TestCheckpointRemoveFromPending(t *testing.T) {
	tmpDir := t.TempDir()

	cm := NewCheckpointManager(tmpDir, 100*time.Millisecond)
	err := cm.Start("wf-1")
	assert.NoError(t, err)
	defer cm.Stop()

	// Set pending tasks
	cm.SetPendingTasks([]string{"task-1", "task-2", "task-3"})

	// Mark task-2 as completed (should remove from pending)
	cm.AddCompletedTask("task-2")

	cp := cm.GetCurrent()
	assert.Equal(t, 2, len(cp.PendingTasks))
	assert.NotContains(t, cp.PendingTasks, "task-2")
	assert.Contains(t, cp.PendingTasks, "task-1")
	assert.Contains(t, cp.PendingTasks, "task-3")
}

func TestCheckpointFailedTask(t *testing.T) {
	tmpDir := t.TempDir()

	cm := NewCheckpointManager(tmpDir, 100*time.Millisecond)
	err := cm.Start("wf-1")
	assert.NoError(t, err)
	defer cm.Stop()

	cm.SetPendingTasks([]string{"task-1", "task-2"})
	cm.AddFailedTask("task-1")

	cp := cm.GetCurrent()
	assert.Equal(t, 1, len(cp.FailedTasks))
	assert.Equal(t, 1, len(cp.PendingTasks))
	assert.Contains(t, cp.FailedTasks, "task-1")
	assert.Contains(t, cp.PendingTasks, "task-2")
}

func TestCheckpointDoubleStart(t *testing.T) {
	tmpDir := t.TempDir()

	cm := NewCheckpointManager(tmpDir, 100*time.Millisecond)
	err := cm.Start("wf-1")
	assert.NoError(t, err)
	defer cm.Stop()

	// Starting again should error
	err = cm.Start("wf-2")
	assert.Error(t, err)
}

func TestCheckpointMultipleStop(t *testing.T) {
	tmpDir := t.TempDir()

	cm := NewCheckpointManager(tmpDir, 100*time.Millisecond)
	cm.Start("wf-1")

	// Multiple stops should not error
	err := cm.Stop()
	assert.NoError(t, err)

	err = cm.Stop()
	assert.NoError(t, err)
}

func TestCheckpointCurrentCopy(t *testing.T) {
	tmpDir := t.TempDir()

	cm := NewCheckpointManager(tmpDir, 100*time.Millisecond)
	cm.Start("wf-1")
	defer cm.Stop()

	cp := cm.GetCurrent()
	// Mutating returned checkpoint shouldn't affect internal state
	cp.Stage = "modified"

	cp2 := cm.GetCurrent()
	assert.NotEqual(t, "modified", cp2.Stage)
}