github.com/yaling888/clash@v1.53.0/common/batch/batch_test.go (about)

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