github.com/yandex/pandora@v0.5.32/core/provider/num_test.go (about)

     1  package provider
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"github.com/stretchr/testify/assert"
     8  	"github.com/yandex/pandora/core"
     9  )
    10  
    11  func Test_Num(t *testing.T) {
    12  	var (
    13  		limit int
    14  
    15  		p      core.Provider
    16  		ctx    context.Context
    17  		cancel context.CancelFunc
    18  		runRes chan error
    19  	)
    20  	beforeEach := func() {
    21  		limit = 0
    22  		runRes = make(chan error)
    23  	}
    24  	justBeforeEach := func() {
    25  		ctx, cancel = context.WithCancel(context.Background())
    26  		p = NewNumConf(NumConfig{limit})
    27  		go func() {
    28  			runRes <- p.Run(ctx, core.ProviderDeps{})
    29  		}()
    30  	}
    31  
    32  	t.Run("unlimited", func(t *testing.T) {
    33  		beforeEach()
    34  		justBeforeEach()
    35  
    36  		for i := 0; i < 100; i++ {
    37  			a, ok := p.Acquire()
    38  			assert.True(t, ok)
    39  			assert.Equal(t, i, a)
    40  		}
    41  		cancel()
    42  
    43  		res := <-runRes
    44  		assert.NoError(t, res)
    45  
    46  		a, ok := p.Acquire()
    47  		assert.False(t, ok)
    48  		assert.Nil(t, a)
    49  	})
    50  
    51  	t.Run("unlimited", func(t *testing.T) {
    52  		beforeEach()
    53  		limit = 50
    54  		justBeforeEach()
    55  
    56  		for i := 0; i < limit; i++ {
    57  			a, ok := p.Acquire()
    58  			assert.True(t, ok)
    59  			assert.Equal(t, i, a)
    60  		}
    61  		res := <-runRes
    62  		assert.NoError(t, res)
    63  
    64  		a, ok := p.Acquire()
    65  		assert.False(t, ok)
    66  		assert.Nil(t, a)
    67  	})
    68  }