github.com/songzhibin97/gkit@v1.2.13/concurrent/stream_test.go (about)

     1  package concurrent
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/stretchr/testify/assert"
     9  )
    10  
    11  func generateStreamNumber(n int) (ret []interface{}) {
    12  	for i := 0; i < n; i++ {
    13  		ret = append(ret, i)
    14  	}
    15  	return ret
    16  }
    17  
    18  var flagSlice = []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
    19  
    20  func filter(i interface{}) bool {
    21  	return i.(int)&1 == 0
    22  }
    23  
    24  func TestStream(t *testing.T) {
    25  	c := Stream(context.Background(), generateStreamNumber(10)...)
    26  	var ret []int
    27  	for v := range c {
    28  		ret = append(ret, v.(int))
    29  	}
    30  	assert.Equal(t, flagSlice, ret)
    31  
    32  	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
    33  	defer cancel()
    34  	c = Stream(ctx, generateStreamNumber(10)...)
    35  	ret = nil
    36  	for v := range c {
    37  		ret = append(ret, v.(int))
    38  		time.Sleep(time.Second)
    39  	}
    40  	assert.Equal(t, []int{0, 1, 2}, ret)
    41  }
    42  
    43  func TestTaskN(t *testing.T) {
    44  	c := Stream(context.Background(), generateStreamNumber(10)...)
    45  	taskN := TaskN(context.Background(), c, 3)
    46  	var ret []int
    47  	for v := range taskN {
    48  		ret = append(ret, v.(int))
    49  	}
    50  	assert.Equal(t, []int{0, 1, 2}, ret)
    51  
    52  	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
    53  	defer cancel()
    54  	c = Stream(context.Background(), generateStreamNumber(10)...)
    55  	taskN = TaskN(ctx, c, 10)
    56  	ret = nil
    57  	for v := range taskN {
    58  		ret = append(ret, v.(int))
    59  		time.Sleep(time.Second)
    60  	}
    61  	assert.Equal(t, []int{0, 1, 2}, ret)
    62  }
    63  
    64  func TestTaskFn(t *testing.T) {
    65  	c := Stream(context.Background(), generateStreamNumber(10)...)
    66  	taskFn := TaskFn(context.Background(), c, filter)
    67  	var ret []int
    68  	for v := range taskFn {
    69  		ret = append(ret, v.(int))
    70  	}
    71  	assert.Equal(t, []int{0, 2, 4, 6, 8}, ret)
    72  
    73  	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
    74  	defer cancel()
    75  	c = Stream(context.Background(), generateStreamNumber(10)...)
    76  	taskFn = TaskFn(ctx, c, filter)
    77  	ret = nil
    78  	for v := range taskFn {
    79  		ret = append(ret, v.(int))
    80  		time.Sleep(time.Second)
    81  	}
    82  	assert.Equal(t, []int{0, 2, 4}, ret)
    83  }
    84  
    85  func TestTaskWhile(t *testing.T) {
    86  	c := Stream(context.Background(), generateStreamNumber(10)...)
    87  	taskWhile := TaskWhile(context.Background(), c, func(v interface{}) bool {
    88  		return v.(int) == 5
    89  	})
    90  	var ret []int
    91  	for v := range taskWhile {
    92  		ret = append(ret, v.(int))
    93  	}
    94  	assert.Equal(t, []int{5}, ret)
    95  
    96  	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
    97  	defer cancel()
    98  	c = Stream(context.Background(), generateStreamNumber(10)...)
    99  	taskWhile = TaskWhile(ctx, c, func(v interface{}) bool {
   100  		return v.(int) == 5
   101  	})
   102  	ret = nil
   103  	for v := range taskWhile {
   104  		ret = append(ret, v.(int))
   105  		time.Sleep(time.Second)
   106  	}
   107  	assert.Equal(t, []int{5}, ret)
   108  }
   109  
   110  func TestSkipN(t *testing.T) {
   111  	c := Stream(context.Background(), generateStreamNumber(10)...)
   112  	skipN := SkipN(context.Background(), c, 3)
   113  	var ret []int
   114  	for v := range skipN {
   115  		ret = append(ret, v.(int))
   116  	}
   117  	assert.Equal(t, []int{3, 4, 5, 6, 7, 8, 9}, ret)
   118  
   119  	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
   120  	defer cancel()
   121  	c = Stream(context.Background(), generateStreamNumber(10)...)
   122  	skipN = SkipN(ctx, c, 3)
   123  	ret = nil
   124  	for v := range skipN {
   125  		ret = append(ret, v.(int))
   126  		time.Sleep(time.Second)
   127  	}
   128  	assert.Equal(t, []int{3, 4, 5}, ret)
   129  }
   130  
   131  func TestSkipFn(t *testing.T) {
   132  	c := Stream(context.Background(), generateStreamNumber(10)...)
   133  	skipFn := SkipFn(context.Background(), c, filter)
   134  	var ret []int
   135  	for v := range skipFn {
   136  		ret = append(ret, v.(int))
   137  	}
   138  	assert.Equal(t, []int{1, 3, 5, 7, 9}, ret)
   139  
   140  	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
   141  	defer cancel()
   142  	c = Stream(context.Background(), generateStreamNumber(10)...)
   143  	skipFn = SkipFn(ctx, c, filter)
   144  	ret = nil
   145  	for v := range skipFn {
   146  		ret = append(ret, v.(int))
   147  		time.Sleep(time.Second)
   148  	}
   149  	assert.Equal(t, []int{1, 3, 5}, ret)
   150  }
   151  
   152  func TestSkipWhile(t *testing.T) {
   153  	c := Stream(context.Background(), generateStreamNumber(10)...)
   154  	skipWhile := SkipWhile(context.Background(), c, func(v interface{}) bool {
   155  		return v.(int) == 0
   156  	})
   157  	var ret []int
   158  	for v := range skipWhile {
   159  		ret = append(ret, v.(int))
   160  	}
   161  	assert.Equal(t, []int{1, 2, 3, 4, 5, 6, 7, 8, 9}, ret)
   162  
   163  	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
   164  	defer cancel()
   165  	c = Stream(context.Background(), generateStreamNumber(10)...)
   166  	skipWhile = SkipWhile(ctx, c, func(v interface{}) bool {
   167  		return v.(int) == 0
   168  	})
   169  	ret = nil
   170  	for v := range skipWhile {
   171  		ret = append(ret, v.(int))
   172  		time.Sleep(time.Second)
   173  	}
   174  	assert.Equal(t, []int{1, 2, 3}, ret)
   175  }