learning/go-patterns/stages/06-bounded/main_test.go

package main

import (
	"fmt"
	"sync"
	"testing"
	"time"
)

func TestRunBoundedHonorsTheConcurrencyLimit(t *testing.T) {
	const (
		taskCount = 6
		limit     = 2
	)

	started := make(chan struct{}, taskCount)
	release := make(chan struct{})
	var mu sync.Mutex
	active := 0
	maxActive := 0
	tasks := make([]Task, taskCount)
	for index := range tasks {
		index := index
		tasks[index] = func() (string, error) {
			mu.Lock()
			active++
			if active > maxActive {
				maxActive = active
			}
			mu.Unlock()
			started <- struct{}{}

			<-release

			mu.Lock()
			active--
			mu.Unlock()
			return fmt.Sprintf("job-%d", index), nil
		}
	}

	done := make(chan struct{})
	var values []string
	var err error
	go func() {
		values, err = RunBounded(tasks, limit)
		close(done)
	}()

	for i := 0; i < limit; i++ {
		select {
		case <-started:
		case <-time.After(time.Second):
			t.Fatalf("timed out waiting for worker %d to start", i+1)
		}
	}
	select {
	case <-started:
		t.Fatal("a third task started while the limit was full")
	default:
	}

	close(release)
	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("RunBounded did not finish after tasks were released")
	}

	if err != nil {
		t.Fatalf("RunBounded returned error: %v", err)
	}
	if len(values) != taskCount {
		t.Fatalf("got %d results, want %d", len(values), taskCount)
	}
	mu.Lock()
	gotMaxActive := maxActive
	mu.Unlock()
	if gotMaxActive > limit {
		t.Fatalf("max active = %d, want at most %d", gotMaxActive, limit)
	}
}

func TestRunBoundedReturnsAnErrorAfterTheWholeBatchFinishes(t *testing.T) {
	const taskCount = 4
	completed := make(chan struct{}, taskCount)
	tasks := make([]Task, taskCount)
	for index := range tasks {
		index := index
		tasks[index] = func() (string, error) {
			completed <- struct{}{}
			if index == 1 {
				return "", fmt.Errorf("job %d failed", index)
			}
			return fmt.Sprintf("job %d", index), nil
		}
	}

	values, err := RunBounded(tasks, 2)
	if err == nil || err.Error() != "job 1 failed" {
		t.Fatalf("error = %v, want job 1 failed", err)
	}
	for index := 0; index < taskCount; index++ {
		select {
		case <-completed:
		case <-time.After(time.Second):
			t.Fatalf("task %d did not complete before RunBounded returned", index)
		}
	}
	if values[0] != "job 0" || values[2] != "job 2" {
		t.Fatalf("values = %#v, want successful results preserved", values)
	}
}