github.com/metacubex/mihomo@v1.18.5/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  		idx := i
    48  		b.Go(strconv.Itoa(idx), func() (string, error) {
    49  			time.Sleep(time.Millisecond * 100)
    50  			return strconv.Itoa(idx), nil
    51  		})
    52  	}
    53  	result, _ := b.WaitAndGetResult()
    54  	duration := time.Since(now)
    55  	assert.Greater(t, duration, time.Millisecond*260)
    56  	assert.Equal(t, 7, len(result))
    57  
    58  	for k, v := range result {
    59  		assert.NoError(t, v.Err)
    60  		assert.Equal(t, k, v.Value)
    61  	}
    62  }
    63  
    64  func TestBatchContext(t *testing.T) {
    65  	b, ctx := New[string](context.Background())
    66  
    67  	b.Go("error", func() (string, error) {
    68  		time.Sleep(time.Millisecond * 100)
    69  		return "", errors.New("test error")
    70  	})
    71  
    72  	b.Go("ctx", func() (string, error) {
    73  		<-ctx.Done()
    74  		return "", ctx.Err()
    75  	})
    76  
    77  	result, err := b.WaitAndGetResult()
    78  
    79  	assert.NotNil(t, err)
    80  	assert.Equal(t, "error", err.Key)
    81  
    82  	assert.Equal(t, ctx.Err(), result["ctx"].Err)
    83  }