picker_test.go 923 B

1234567891011121314151617181920212223242526272829303132333435363738394041
  1. package picker
  2. import (
  3. "context"
  4. "testing"
  5. "time"
  6. "github.com/samber/lo"
  7. "github.com/stretchr/testify/assert"
  8. )
  9. func sleepAndSend[T any](ctx context.Context, delay int, input T) func() (T, error) {
  10. return func() (T, error) {
  11. timer := time.NewTimer(time.Millisecond * time.Duration(delay))
  12. select {
  13. case <-timer.C:
  14. return input, nil
  15. case <-ctx.Done():
  16. return lo.Empty[T](), ctx.Err()
  17. }
  18. }
  19. }
  20. func TestPicker_Basic(t *testing.T) {
  21. picker, ctx := WithContext[int](context.Background())
  22. picker.Go(sleepAndSend(ctx, 30, 2))
  23. picker.Go(sleepAndSend(ctx, 20, 1))
  24. number := picker.Wait()
  25. assert.NotNil(t, number)
  26. assert.Equal(t, number, 1)
  27. }
  28. func TestPicker_Timeout(t *testing.T) {
  29. picker, ctx := WithTimeout[int](context.Background(), time.Millisecond*5)
  30. picker.Go(sleepAndSend(ctx, 20, 1))
  31. number := picker.Wait()
  32. assert.Equal(t, number, lo.Empty[int]())
  33. assert.NotNil(t, picker.Error())
  34. }