batch_test.go 1.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283
  1. package batch
  2. import (
  3. "context"
  4. "errors"
  5. "strconv"
  6. "testing"
  7. "time"
  8. "github.com/stretchr/testify/assert"
  9. )
  10. func TestBatch(t *testing.T) {
  11. b, _ := New[string](context.Background())
  12. now := time.Now()
  13. b.Go("foo", func() (string, error) {
  14. time.Sleep(time.Millisecond * 100)
  15. return "foo", nil
  16. })
  17. b.Go("bar", func() (string, error) {
  18. time.Sleep(time.Millisecond * 150)
  19. return "bar", nil
  20. })
  21. result, err := b.WaitAndGetResult()
  22. assert.Nil(t, err)
  23. duration := time.Since(now)
  24. assert.Less(t, duration, time.Millisecond*200)
  25. assert.Equal(t, 2, len(result))
  26. for k, v := range result {
  27. assert.NoError(t, v.Err)
  28. assert.Equal(t, k, v.Value)
  29. }
  30. }
  31. func TestBatchWithConcurrencyNum(t *testing.T) {
  32. b, _ := New[string](
  33. context.Background(),
  34. WithConcurrencyNum[string](3),
  35. )
  36. now := time.Now()
  37. for i := 0; i < 7; i++ {
  38. idx := i
  39. b.Go(strconv.Itoa(idx), func() (string, error) {
  40. time.Sleep(time.Millisecond * 100)
  41. return strconv.Itoa(idx), nil
  42. })
  43. }
  44. result, _ := b.WaitAndGetResult()
  45. duration := time.Since(now)
  46. assert.Greater(t, duration, time.Millisecond*260)
  47. assert.Equal(t, 7, len(result))
  48. for k, v := range result {
  49. assert.NoError(t, v.Err)
  50. assert.Equal(t, k, v.Value)
  51. }
  52. }
  53. func TestBatchContext(t *testing.T) {
  54. b, ctx := New[string](context.Background())
  55. b.Go("error", func() (string, error) {
  56. time.Sleep(time.Millisecond * 100)
  57. return "", errors.New("test error")
  58. })
  59. b.Go("ctx", func() (string, error) {
  60. <-ctx.Done()
  61. return "", ctx.Err()
  62. })
  63. result, err := b.WaitAndGetResult()
  64. assert.NotNil(t, err)
  65. assert.Equal(t, "error", err.Key)
  66. assert.Equal(t, ctx.Err(), result["ctx"].Err)
  67. }