add tests and cleanup when context stop

Author: Codelax

internal/tasks/tasks.go

@@ -1,6 +1,7 @@
 package tasks
 
 import (
+	"context"
 	"fmt"
 	"time"
 
@@ -70,7 +71,7 @@ func (ts *Tasks) Cleanup(failed int) {
 }
 
 // Execute tasks with interactive display and cleanup on fail
-func (ts *Tasks) Execute(data interface{}) (interface{}, error) {
+func (ts *Tasks) Execute(ctx context.Context, data interface{}) (interface{}, error) {
 	var err error
 	totalTasks := len(ts.tasks)
 	spin := spinner.New(spinner.CharSets[11], 100*time.Millisecond)
@@ -88,6 +89,15 @@ func (ts *Tasks) Execute(data interface{}) (interface{}, error) {
 			return nil, fmt.Errorf("task %d %q failed: %w", i+1, task.Name, err)
 		}
 
+		select {
+		case <-ctx.Done():
+			spin.Stop()
+			fmt.Println("context canceled, cleaning up created resources")
+			ts.Cleanup(i + 1)
+			return nil, fmt.Errorf("task %d %q failed: context canceled", i+1, task.Name)
+		default:
+		}
+
 		spin.Stop()
 	}
 

internal/tasks/tasks_test.go

@@ -0,0 +1,68 @@
+package tasks_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/alecthomas/assert"
+	"github.com/scaleway/scaleway-cli/v2/internal/tasks"
+)
+
+func TestCleanup(t *testing.T) {
+	ts := tasks.Begin()
+
+	clean := 0
+
+	tasks.AddWithCleanUp(ts, "Task 1", func(_ interface{}) (interface{}, string, error) {
+		return nil, "", nil
+	}, func(string) error {
+		clean += 1
+		return nil
+	})
+	tasks.AddWithCleanUp(ts, "Task 2", func(_ interface{}) (interface{}, string, error) {
+		return nil, "", nil
+	}, func(string) error {
+		clean += 1
+		return nil
+	})
+	tasks.AddWithCleanUp(ts, "Task 3", func(_ interface{}) (interface{}, string, error) {
+		return nil, "", fmt.Errorf("fail")
+	}, func(string) error {
+		clean += 1
+		return nil
+	})
+	_, err := ts.Execute(context.Background(), nil)
+	assert.NotNil(t, err, "Execute should return error after cleanup")
+	assert.Equal(t, clean, 2, "2 task cleanup should have been executed")
+}
+
+func TestCleanupOnContext(t *testing.T) {
+	ts := tasks.Begin()
+
+	clean := 0
+	ctx, cancel := context.WithCancel(context.Background())
+
+	tasks.AddWithCleanUp(ts, "Task 1", func(_ interface{}) (interface{}, string, error) {
+		return nil, "", nil
+	}, func(string) error {
+		clean += 1
+		return nil
+	})
+	tasks.AddWithCleanUp(ts, "Task 2", func(_ interface{}) (interface{}, string, error) {
+		return nil, "", nil
+	}, func(string) error {
+		clean += 1
+		return nil
+	})
+	tasks.AddWithCleanUp(ts, "Task 3", func(_ interface{}) (interface{}, string, error) {
+		cancel()
+		return nil, "", nil
+	}, func(string) error {
+		clean += 1
+		return nil
+	})
+	_, err := ts.Execute(ctx, nil)
+	assert.NotNil(t, err)
+	assert.Equal(t, clean, 3, "3 task cleanup should have been executed")
+}