github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/retry/budget/budget_test.go (about)

     1  package budget
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/jonboulle/clockwork"
     9  	"github.com/stretchr/testify/require"
    10  
    11  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
    12  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
    13  )
    14  
    15  func TestUnlimitedBudget(t *testing.T) {
    16  	xtest.TestManyTimes(t, func(t testing.TB) {
    17  		ctx, cancel := xcontext.WithCancel(xtest.Context(t))
    18  		q := Limited(-1)
    19  		require.NoError(t, q.Acquire(ctx))
    20  		cancel()
    21  		require.ErrorIs(t, q.Acquire(ctx), context.Canceled)
    22  	})
    23  }
    24  
    25  func TestLimited(t *testing.T) {
    26  	xtest.TestManyTimes(t, func(t testing.TB) {
    27  		ctx, cancel := xcontext.WithCancel(xtest.Context(t))
    28  		clock := clockwork.NewFakeClock()
    29  		q := Limited(1, withFixedBudgetClock(clock))
    30  		defer q.Stop()
    31  		require.NoError(t, q.Acquire(ctx))
    32  		acquireCh := make(chan struct{})
    33  		go func() {
    34  			err := q.Acquire(ctx)
    35  			acquireCh <- struct{}{}
    36  			require.NoError(t, err)
    37  		}()
    38  		timeCh := make(chan struct{})
    39  		go func() {
    40  			clock.Advance(time.Second - time.Nanosecond)
    41  			timeCh <- struct{}{}
    42  			clock.Advance(time.Nanosecond)
    43  		}()
    44  		select {
    45  		case <-acquireCh:
    46  			require.Fail(t, "")
    47  		case <-timeCh:
    48  		}
    49  		<-acquireCh
    50  		cancel()
    51  		require.ErrorIs(t, q.Acquire(ctx), context.Canceled)
    52  	})
    53  }
    54  
    55  func TestPercent(t *testing.T) {
    56  	xtest.TestManyTimes(t, func(t testing.TB) {
    57  		var (
    58  			total   = 1000000
    59  			percent = 0.25
    60  			ctx     = xtest.Context(t)
    61  			b       = Percent(int(percent * 100))
    62  			success int
    63  		)
    64  		for i := 0; i < total; i++ {
    65  			if b.Acquire(ctx) == nil {
    66  				success++
    67  			}
    68  		}
    69  		require.GreaterOrEqual(t, success, int(float64(total)*(percent-0.1*percent)))
    70  		require.LessOrEqual(t, success, int(float64(total)*(percent+0.1*percent)))
    71  	}, xtest.StopAfter(5*time.Second))
    72  }