fix test

Signed-off-by: Cabinfever_B <cabinfeveroier@gmail.com>
tikv · Jan 17, 2023 · 122c52b · 122c52b
1 parent 0c46eaf
commit 122c52b
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/pkg/mcs/resource_manager/client/limiter.go b/pkg/mcs/resource_manager/client/limiter.go
@@ -84,10 +84,10 @@ func (lim *Limiter) Limit() Limit {
 
 // NewLimiter returns a new Limiter that allows events up to rate r and permits
 // bursts of at most b tokens.
-func NewLimiter(r Limit, tokens, maxRequestTokens float64, lowTokensNotifyChan chan struct{}) *Limiter {
+func NewLimiter(now time.Time, r Limit, tokens, maxRequestTokens float64, lowTokensNotifyChan chan struct{}) *Limiter {
 	lim := &Limiter{
 		limit:               r,
-		last:                time.Now(),
+		last:                now,
 		tokens:              tokens,
 		maxRequestTokens:    maxRequestTokens,
 		lowTokensNotifyChan: lowTokensNotifyChan,

diff --git a/pkg/mcs/resource_manager/client/limiter_test.go b/pkg/mcs/resource_manager/client/limiter_test.go
@@ -79,11 +79,11 @@ func runReserve(t *testing.T, lim *Limiter, req request, maxReserve time.Duratio
 }
 
 func checkTokens(re *require.Assertions, lim *Limiter, t time.Time, expected float64) {
-	re.LessOrEqual(math.Abs(expected-lim.AvailableTokens(t)), 1e-2)
+	re.LessOrEqual(math.Abs(expected-lim.AvailableTokens(t)), 1e-7)
 }
 
 func TestSimpleReserve(t *testing.T) {
-	lim := NewLimiter(1, 2, 1000, make(chan struct{}, 1))
+	lim := NewLimiter(t0, 1, 2, 1000, make(chan struct{}, 1))
 
 	runReserveMax(t, lim, request{t0, 3, t1, true})
 	runReserveMax(t, lim, request{t0, 3, t4, true})
@@ -97,7 +97,7 @@ func TestSimpleReserve(t *testing.T) {
 
 func TestReconfig(t *testing.T) {
 	re := require.New(t)
-	lim := NewLimiter(1, 2, 1000, make(chan struct{}, 1))
+	lim := NewLimiter(t0, 1, 2, 1000, make(chan struct{}, 1))
 
 	runReserveMax(t, lim, request{t0, 4, t2, true})
 	args := tokenBucketReconfigureArgs{
@@ -111,7 +111,7 @@ func TestReconfig(t *testing.T) {
 
 func TestNotify(t *testing.T) {
 	nc := make(chan struct{}, 1)
-	lim := NewLimiter(1, 0, 1000, nc)
+	lim := NewLimiter(t0, 1, 0, 1000, nc)
 
 	args := tokenBucketReconfigureArgs{
 		NewTokens:       1000.,
@@ -132,8 +132,8 @@ func TestCancel(t *testing.T) {
 	ctx1, cancel1 := context.WithDeadline(ctx, t2)
 	re := require.New(t)
 	nc := make(chan struct{}, 1)
-	lim1 := NewLimiter(1, 10, 100, nc)
-	lim2 := NewLimiter(1, 0, 100, nc)
+	lim1 := NewLimiter(t0, 1, 10, 100, nc)
+	lim2 := NewLimiter(t0, 1, 0, 100, nc)
 
 	r1 := runReserveMax(t, lim1, request{t0, 5, t0, true})
 	checkTokens(re, lim1, t0, 5)