github.com/lingyao2333/mo-zero@v1.4.1/core/collection/rollingwindow_test.go (about)

     1  package collection
     2  
     3  import (
     4  	"math/rand"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/lingyao2333/mo-zero/core/stringx"
     9  	"github.com/stretchr/testify/assert"
    10  )
    11  
    12  const duration = time.Millisecond * 50
    13  
    14  func TestNewRollingWindow(t *testing.T) {
    15  	assert.NotNil(t, NewRollingWindow(10, time.Second))
    16  	assert.Panics(t, func() {
    17  		NewRollingWindow(0, time.Second)
    18  	})
    19  }
    20  
    21  func TestRollingWindowAdd(t *testing.T) {
    22  	const size = 3
    23  	r := NewRollingWindow(size, duration)
    24  	listBuckets := func() []float64 {
    25  		var buckets []float64
    26  		r.Reduce(func(b *Bucket) {
    27  			buckets = append(buckets, b.Sum)
    28  		})
    29  		return buckets
    30  	}
    31  	assert.Equal(t, []float64{0, 0, 0}, listBuckets())
    32  	r.Add(1)
    33  	assert.Equal(t, []float64{0, 0, 1}, listBuckets())
    34  	elapse()
    35  	r.Add(2)
    36  	r.Add(3)
    37  	assert.Equal(t, []float64{0, 1, 5}, listBuckets())
    38  	elapse()
    39  	r.Add(4)
    40  	r.Add(5)
    41  	r.Add(6)
    42  	assert.Equal(t, []float64{1, 5, 15}, listBuckets())
    43  	elapse()
    44  	r.Add(7)
    45  	assert.Equal(t, []float64{5, 15, 7}, listBuckets())
    46  }
    47  
    48  func TestRollingWindowReset(t *testing.T) {
    49  	const size = 3
    50  	r := NewRollingWindow(size, duration, IgnoreCurrentBucket())
    51  	listBuckets := func() []float64 {
    52  		var buckets []float64
    53  		r.Reduce(func(b *Bucket) {
    54  			buckets = append(buckets, b.Sum)
    55  		})
    56  		return buckets
    57  	}
    58  	r.Add(1)
    59  	elapse()
    60  	assert.Equal(t, []float64{0, 1}, listBuckets())
    61  	elapse()
    62  	assert.Equal(t, []float64{1}, listBuckets())
    63  	elapse()
    64  	assert.Nil(t, listBuckets())
    65  
    66  	// cross window
    67  	r.Add(1)
    68  	time.Sleep(duration * 10)
    69  	assert.Nil(t, listBuckets())
    70  }
    71  
    72  func TestRollingWindowReduce(t *testing.T) {
    73  	const size = 4
    74  	tests := []struct {
    75  		win    *RollingWindow
    76  		expect float64
    77  	}{
    78  		{
    79  			win:    NewRollingWindow(size, duration),
    80  			expect: 10,
    81  		},
    82  		{
    83  			win:    NewRollingWindow(size, duration, IgnoreCurrentBucket()),
    84  			expect: 4,
    85  		},
    86  	}
    87  
    88  	for _, test := range tests {
    89  		t.Run(stringx.Rand(), func(t *testing.T) {
    90  			r := test.win
    91  			for x := 0; x < size; x++ {
    92  				for i := 0; i <= x; i++ {
    93  					r.Add(float64(i))
    94  				}
    95  				if x < size-1 {
    96  					elapse()
    97  				}
    98  			}
    99  			var result float64
   100  			r.Reduce(func(b *Bucket) {
   101  				result += b.Sum
   102  			})
   103  			assert.Equal(t, test.expect, result)
   104  		})
   105  	}
   106  }
   107  
   108  func TestRollingWindowBucketTimeBoundary(t *testing.T) {
   109  	const size = 3
   110  	interval := time.Millisecond * 30
   111  	r := NewRollingWindow(size, interval)
   112  	listBuckets := func() []float64 {
   113  		var buckets []float64
   114  		r.Reduce(func(b *Bucket) {
   115  			buckets = append(buckets, b.Sum)
   116  		})
   117  		return buckets
   118  	}
   119  	assert.Equal(t, []float64{0, 0, 0}, listBuckets())
   120  	r.Add(1)
   121  	assert.Equal(t, []float64{0, 0, 1}, listBuckets())
   122  	time.Sleep(time.Millisecond * 45)
   123  	r.Add(2)
   124  	r.Add(3)
   125  	assert.Equal(t, []float64{0, 1, 5}, listBuckets())
   126  	// sleep time should be less than interval, and make the bucket change happen
   127  	time.Sleep(time.Millisecond * 20)
   128  	r.Add(4)
   129  	r.Add(5)
   130  	r.Add(6)
   131  	assert.Equal(t, []float64{1, 5, 15}, listBuckets())
   132  	time.Sleep(time.Millisecond * 100)
   133  	r.Add(7)
   134  	r.Add(8)
   135  	r.Add(9)
   136  	assert.Equal(t, []float64{0, 0, 24}, listBuckets())
   137  }
   138  
   139  func TestRollingWindowDataRace(t *testing.T) {
   140  	const size = 3
   141  	r := NewRollingWindow(size, duration)
   142  	stop := make(chan bool)
   143  	go func() {
   144  		for {
   145  			select {
   146  			case <-stop:
   147  				return
   148  			default:
   149  				r.Add(float64(rand.Int63()))
   150  				time.Sleep(duration / 2)
   151  			}
   152  		}
   153  	}()
   154  	go func() {
   155  		for {
   156  			select {
   157  			case <-stop:
   158  				return
   159  			default:
   160  				r.Reduce(func(b *Bucket) {})
   161  			}
   162  		}
   163  	}()
   164  	time.Sleep(duration * 5)
   165  	close(stop)
   166  }
   167  
   168  func elapse() {
   169  	time.Sleep(duration)
   170  }