go-study/routine/pool_test.go

package routine

import (
	"fmt"
	"runtime"
	"runtime/debug"
	"sync/atomic"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestPool(t *testing.T) {
	num := runtime.NumCPU()
	var sum atomic.Int32

	task := func(t T) {
		num := t.(int32)
		if num < 0 {
			panic("unable to handle negative numbers")
		}

		sum.Add(num)
	}
	handler := func(r any) {
		fmt.Printf("Panic: %v\n %s", r, string(debug.Stack()))
	}
	pool := NewPool(
		WithWorkers(num),
		WithCapacity(num),
		WithTaskFn(task),
		WithPanicHandler(handler),
	)
	pool.Start()

	for i := int32(1000); i >= -1; i-- {
		pool.Push(i)
	}

	pool.Wait()

	if assert.Equal(t, int32(500500), sum.Load()) {
		t.Log("the sum value is right")
	}
}