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