From d7a6409ed1f332070fc97bf3928bf15725c06f04 Mon Sep 17 00:00:00 2001
From: Neur0toxine <pashok9825@gmail.com>
Date: Fri, 16 Feb 2024 18:10:34 +0300
Subject: [PATCH] fix data race in test

---
 v1/rate_limit_test.go | 23 +++++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/v1/rate_limit_test.go b/v1/rate_limit_test.go
index c845705..7e9ab72 100644
--- a/v1/rate_limit_test.go
+++ b/v1/rate_limit_test.go
@@ -3,6 +3,7 @@ package v1
 import (
 	"github.com/jonboulle/clockwork"
 	"github.com/stretchr/testify/suite"
+	"runtime"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -21,8 +22,20 @@ func (t *TokensBucketTest) Test_NewTokensBucket() {
 	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() {
-	tb := NewTokensBucket(100, time.Hour, time.Minute)
+	tb := t.new(100, time.Hour, time.Minute, &realSleeper{})
 	start := time.Now()
 	for i := 0; i < 100; i++ {
 		tb.Obtain("a")
@@ -32,9 +45,7 @@ func (t *TokensBucketTest) Test_Obtain_NoThrottle() {
 
 func (t *TokensBucketTest) Test_Obtain_Sleep() {
 	clock := &fakeSleeper{}
-	tb := NewTokensBucket(100, time.Hour, time.Minute)
-	tb.cancel.Store(true) // prevent unused token removal.
-	tb.sleep = clock
+	tb := t.new(100, time.Hour, time.Minute, clock)
 
 	var wg sync.WaitGroup
 	wg.Add(1)
@@ -51,8 +62,8 @@ func (t *TokensBucketTest) Test_Obtain_Sleep() {
 
 func (t *TokensBucketTest) Test_Obtain_AddRPS() {
 	clock := clockwork.NewFakeClock()
-	tb := NewTokensBucket(100, time.Hour, time.Minute)
-	tb.sleep = clock
+	tb := t.new(100, time.Hour, time.Minute, clock)
+	go tb.deleteUnusedToken()
 	tb.Obtain("a")
 	clock.Advance(time.Minute * 2)