learning-go-for-sre-interview/02-worker-pool/main_test.go

package main

import (
	"sort"
	"sync/atomic"
	"testing"
	"time"
)

func TestAllJobsProcessedAndOutputCloses(t *testing.T) {
	jobs := make(chan int)
	go func() {
		for i := 0; i < 100; i++ {
			jobs <- i
		}
		close(jobs)
	}()

	out := Pool(jobs, 5, func(n int) int { return n * 2 })

	got := []int{}
	done := make(chan struct{})
	go func() {
		for v := range out { // terminates only if out is closed
			got = append(got, v)
		}
		close(done)
	}()

	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("output channel never closed — deadlock or missing close(out)")
	}

	if len(got) != 100 {
		t.Fatalf("want 100 results, got %d", len(got))
	}
	sort.Ints(got)
	for i, v := range got {
		if v != i*2 {
			t.Fatalf("missing/wrong result: index %d = %d", i, v)
		}
	}
}

func TestActuallyConcurrent(t *testing.T) {
	var inFlight, peak atomic.Int32
	jobs := make(chan int)
	go func() {
		for i := 0; i < 12; i++ {
			jobs <- i
		}
		close(jobs)
	}()

	out := Pool(jobs, 4, func(n int) int {
		cur := inFlight.Add(1)
		for {
			p := peak.Load()
			if cur <= p || peak.CompareAndSwap(p, cur) {
				break
			}
		}
		time.Sleep(30 * time.Millisecond)
		inFlight.Add(-1)
		return n
	})
	drained := make(chan struct{})
	go func() {
		for range out {
		}
		close(drained)
	}()
	select {
	case <-drained:
	case <-time.After(3 * time.Second):
		t.Fatal("output never closed")
	}

	if p := peak.Load(); p < 2 {
		t.Errorf("expected concurrent execution, peak parallelism was %d", p)
	} else if p > 4 {
		t.Errorf("more than `workers` goroutines running fn: peak %d > 4", p)
	}
}

func TestEmptyJobs(t *testing.T) {
	jobs := make(chan int)
	close(jobs)
	out := Pool(jobs, 3, func(n int) int { return n })
	select {
	case _, ok := <-out:
		if ok {
			t.Fatal("got a value from empty job stream")
		}
	case <-time.After(time.Second):
		t.Fatal("output not closed for empty job stream")
	}
}