library-go/worker_pool/worker_pool_test.go

90 lines
1.7 KiB
Go
Raw Permalink Normal View History

2024-08-20 15:33:20 +00:00
package workerpool
2024-08-20 15:11:32 +00:00
import (
"errors"
"sync"
"testing"
"time"
2024-08-20 15:33:20 +00:00
"github.com/stretchr/testify/assert"
2024-08-20 15:11:32 +00:00
)
func TestNewWorkerPool(t *testing.T) {
t.Run("default size pool", func(t *testing.T) {
pool := NewWorkerPool(0)
assert.NotNil(t, pool)
})
t.Run("custom size pool", func(t *testing.T) {
size := 100
pool := NewWorkerPool(size)
assert.NotNil(t, pool)
})
}
func TestSubmit(t *testing.T) {
t.Run("submit task to worker pool", func(t *testing.T) {
pool := NewWorkerPool(10)
var wg sync.WaitGroup
wg.Add(1)
err := pool.Submit(func() {
defer wg.Done()
time.Sleep(100 * time.Millisecond)
})
assert.NoError(t, err)
wg.Wait()
})
}
func TestSubmitAndWaitAll(t *testing.T) {
t.Run("submit and wait all tasks succeed", func(t *testing.T) {
pool := NewWorkerPool(10)
tasks := []func() error{
func() error {
time.Sleep(100 * time.Millisecond)
return nil
},
func() error {
time.Sleep(50 * time.Millisecond)
return nil
},
}
taskErrCh, submitErr := pool.SubmitAndWaitAll(tasks...)
assert.NoError(t, submitErr)
close(taskErrCh)
for err := range taskErrCh {
assert.NoError(t, err)
}
})
t.Run("submit and wait all tasks with errors", func(t *testing.T) {
pool := NewWorkerPool(10)
expectedError := errors.New("task error")
tasks := []func() error{
func() error {
time.Sleep(100 * time.Millisecond)
return nil
},
func() error {
time.Sleep(50 * time.Millisecond)
return expectedError
},
}
taskErrCh, submitErr := pool.SubmitAndWaitAll(tasks...)
assert.NoError(t, submitErr)
close(taskErrCh)
foundError := false
for err := range taskErrCh {
if err != nil {
foundError = true
assert.Equal(t, expectedError, err)
}
}
assert.True(t, foundError)
})
}