github.com/metacubex/mihomo@v1.18.5/common/picker/picker_test.go (about)

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