1
0
mirror of synced 2024-11-24 05:56:04 +03:00

fix data race in test

This commit is contained in:
Pavel 2024-02-16 18:10:34 +03:00
parent f3f87d1b7b
commit d7a6409ed1

View File

@ -3,6 +3,7 @@ package v1
import ( import (
"github.com/jonboulle/clockwork" "github.com/jonboulle/clockwork"
"github.com/stretchr/testify/suite" "github.com/stretchr/testify/suite"
"runtime"
"sync" "sync"
"sync/atomic" "sync/atomic"
"testing" "testing"
@ -21,8 +22,20 @@ func (t *TokensBucketTest) Test_NewTokensBucket() {
t.Assert().NotNil(NewTokensBucket(10, time.Hour, time.Hour)) t.Assert().NotNil(NewTokensBucket(10, time.Hour, time.Hour))
} }
func (t *TokensBucketTest) new(
maxRPS uint32, unusedTokenTime, checkTokenTime time.Duration, sleeper sleeper) *TokensBucket {
bucket := &TokensBucket{
maxRPS: maxRPS,
unusedTokenTime: unusedTokenTime,
checkTokenTime: checkTokenTime,
sleep: sleeper,
}
runtime.SetFinalizer(bucket, destructBasket)
return bucket
}
func (t *TokensBucketTest) Test_Obtain_NoThrottle() { func (t *TokensBucketTest) Test_Obtain_NoThrottle() {
tb := NewTokensBucket(100, time.Hour, time.Minute) tb := t.new(100, time.Hour, time.Minute, &realSleeper{})
start := time.Now() start := time.Now()
for i := 0; i < 100; i++ { for i := 0; i < 100; i++ {
tb.Obtain("a") tb.Obtain("a")
@ -32,9 +45,7 @@ func (t *TokensBucketTest) Test_Obtain_NoThrottle() {
func (t *TokensBucketTest) Test_Obtain_Sleep() { func (t *TokensBucketTest) Test_Obtain_Sleep() {
clock := &fakeSleeper{} clock := &fakeSleeper{}
tb := NewTokensBucket(100, time.Hour, time.Minute) tb := t.new(100, time.Hour, time.Minute, clock)
tb.cancel.Store(true) // prevent unused token removal.
tb.sleep = clock
var wg sync.WaitGroup var wg sync.WaitGroup
wg.Add(1) wg.Add(1)
@ -51,8 +62,8 @@ func (t *TokensBucketTest) Test_Obtain_Sleep() {
func (t *TokensBucketTest) Test_Obtain_AddRPS() { func (t *TokensBucketTest) Test_Obtain_AddRPS() {
clock := clockwork.NewFakeClock() clock := clockwork.NewFakeClock()
tb := NewTokensBucket(100, time.Hour, time.Minute) tb := t.new(100, time.Hour, time.Minute, clock)
tb.sleep = clock go tb.deleteUnusedToken()
tb.Obtain("a") tb.Obtain("a")
clock.Advance(time.Minute * 2) clock.Advance(time.Minute * 2)