learning/go-patterns/stages/07-cancel-join/main_test.go

package main

import (
	"context"
	"errors"
	"testing"
	"time"
)

func waitFor(t *testing.T, signal <-chan struct{}, name string) {
	t.Helper()
	select {
	case <-signal:
	case <-time.After(time.Second):
		t.Fatalf("timed out waiting for %s", name)
	}
}

func TestRunFailFastCancelsAndJoinsCooperativeTasks(t *testing.T) {
	wantErr := errors.New("first task failed")
	firstStarted := make(chan struct{})
	waitingStarted := make(chan struct{})
	releaseFailure := make(chan struct{})
	waitingStopped := make(chan struct{})

	tasks := []Task{
		func(context.Context) (string, error) {
			close(firstStarted)
			<-releaseFailure
			return "", wantErr
		},
		func(ctx context.Context) (string, error) {
			close(waitingStarted)
			<-ctx.Done()
			close(waitingStopped)
			return "", ctx.Err()
		},
	}

	done := make(chan struct{})
	var err error
	go func() {
		_, err = RunFailFast(context.Background(), tasks, 2)
		close(done)
	}()

	waitFor(t, firstStarted, "failing task to start")
	waitFor(t, waitingStarted, "cancellable task to start")
	close(releaseFailure)
	waitFor(t, done, "RunFailFast to return")

	if !errors.Is(err, wantErr) {
		t.Fatalf("error = %v, want %v", err, wantErr)
	}
	select {
	case <-waitingStopped:
	default:
		t.Fatal("RunFailFast returned before the cancelled task stopped")
	}
}

func TestRunFailFastReturnsSuccessfulResultsThatFinishedBeforeCancellation(t *testing.T) {
	wantErr := errors.New("request failed")
	finished := make(chan struct{})
	values, err := RunFailFast(context.Background(), []Task{
		func(context.Context) (string, error) { close(finished); return "already done", nil },
		func(context.Context) (string, error) { <-finished; return "", wantErr },
	}, 2)

	if !errors.Is(err, wantErr) {
		t.Fatalf("error = %v, want %v", err, wantErr)
	}
	if values[0] != "already done" {
		t.Fatalf("values = %#v, want successful result to be retained", values)
	}
}