learning/go-patterns/patterns/cancellation/cancellation_test.go

package cancellation

import (
	"context"
	"errors"
	"reflect"
	"testing"
)

func TestProcessStopsBeforeDoingWorkWhenCancelled(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	called := false
	jobs := make(chan int, 1)
	jobs <- 1
	got, err := Process(ctx, jobs, func(value int) int {
		called = true
		return value
	})

	if !errors.Is(err, context.Canceled) {
		t.Fatalf("error = %v, want context.Canceled", err)
	}
	if called {
		t.Fatal("Process() called fn after cancellation")
	}
	if got != nil {
		t.Fatalf("results = %#v, want nil", got)
	}
}

func TestProcessReturnsCompletedResultsOnCancellation(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	jobs := make(chan int, 3)
	jobs <- 1
	jobs <- 2
	jobs <- 3
	close(jobs)

	got, err := Process(ctx, jobs, func(value int) int {
		if value == 2 {
			cancel()
		}
		return value * 10
	})

	if !errors.Is(err, context.Canceled) {
		t.Fatalf("error = %v, want context.Canceled", err)
	}
	if want := []int{10, 20}; !reflect.DeepEqual(got, want) {
		t.Fatalf("results = %#v, want %#v", got, want)
	}
}

func TestProcessConsumesAllJobsWhenChannelCloses(t *testing.T) {
	jobs := make(chan int, 3)
	jobs <- 2
	jobs <- 4
	jobs <- 6
	close(jobs)

	got, err := Process(context.Background(), jobs, func(value int) int { return value / 2 })
	if err != nil {
		t.Fatalf("Process() error = %v", err)
	}
	if want := []int{1, 2, 3}; !reflect.DeepEqual(got, want) {
		t.Fatalf("results = %#v, want %#v", got, want)
	}
}