func TestLowReplenishRate(t *testing.T) {
	rdb := redis.NewClient(&redis.Options{
		Addr: "localhost:6379",
	})
	keys := makeKeys()
	mustInitKeys(t, rdb, keys, 0, TestBaseTime)
	const testReplenishRate = 0.5
	argTimeSec, argTimeUs := timeToRedisPair(TestBaseTime.Add(1 * time.Second))
	resp := updateTokenBucketScript.Eval(
		rdb,
		keys.AsSlice(),
		testReplenishRate,
		TestCapacity,
		1,
		argTimeSec,
		argTimeUs,
		"true",
	)
	assertRedisResults(t, resp,
		false,
		0,
		argTimeSec-1,
		argTimeUs,
	)
	argTimeSec, argTimeUs = timeToRedisPair(TestBaseTime.Add(2 * time.Second))
	resp = updateTokenBucketScript.Eval(
		rdb,
		keys.AsSlice(),
		testReplenishRate,
		TestCapacity,
		1,
		argTimeSec,
		argTimeUs,
		"true",
	)
	assertRedisResults(t, resp,
		true,
		0,
		argTimeSec,
		argTimeUs,
	)
}
Ensure that the token bucket works with very low (i.e. less than one
token/second) replenishment rates.
package redistb
import (
	"context"
	"crypto/rand"
	"math/big"
	"testing"
	"time"
	"github.com/go-redis/redis"
	"github.com/google/trillian/util/clock"
)
const (
	TestCapacity = 5
	TestReplenishRate = 2
	TestBaseTimeSec = 123
	TestBaseTimeUs = 100
	MicrosecondsInSecond = 1000000
	SingleTokenDripTimeUs = MicrosecondsInSecond / TestReplenishRate
	TTL = 60
)
var TestBaseTime = time.Unix(TestBaseTimeSec, TestBaseTimeUs*int64(time.Microsecond/time.Nanosecond))
func TestRedisTokenBucketIntegration(t *testing.T) {
	ctx := context.Background()
	start := time.Now()
	fixedTimeSource := clock.NewFake(start)
	rdb := redis.NewClient(&redis.Options{
		Addr: "localhost:6379",
	})
	tb := New(rdb)
	tb.testing = true
	tb.timeSource = fixedTimeSource
	if err := tb.Load(ctx); err != nil {
		t.Fatalf("failed to load script: %v", err)
	}
	prefix := stringWithCharset(20, alphanumeric)
	call := func(number int) (bool, int64) {
		allowed, remaining, err := tb.Call(
			ctx,
			prefix,
			TestCapacity,
			TestReplenishRate,
			number,
		)
		if err != nil {
			t.Fatal(err)
		}
		return allowed, remaining
	}
	for i := int64(0); i < TestCapacity; i++ {
		allowed, remaining := call(1)
		if !allowed {
			t.Fatal("expected to be allowed")
		}
		if expected := TestCapacity - i - 1; remaining != expected {
			t.Fatalf("expected remaining to be %d, got %d", expected, remaining)
		}
	}
	allowed, remaining := call(1)
	if allowed {
		t.Fatal("expected to be denied")
	}
	if remaining != 0 {
		t.Fatalf("expected remaining to be 0, got %d", remaining)
	}
	singleTokenReplenishTime := 1.0 / TestReplenishRate
	for i := 0; i < 5; i++ {
		timeStep := float64(i+1) * singleTokenReplenishTime
		justBefore := start
		justBefore = justBefore.Add(time.Duration(timeStep * float64(time.Second)))
		justBefore = justBefore.Add(-1 * time.Second / 100)
		fixedTimeSource.Set(justBefore)
		allowed, remaining := call(1)
		if allowed {
			t.Fatalf("expected request before reaching replenish time to be denied on iteration %d", i)
		}
		if remaining != 0 {
			t.Fatalf("expected remaining to be 0, got %d", remaining)
		}
		fixedTimeSource.Set(start.Add(time.Duration(timeStep * float64(time.Second))))
		allowed, remaining = call(1)
		if allowed {
			t.Fatalf("expected first request to be allowed on iteration %d", i)
		}
		if remaining != 0 {
			t.Fatalf("expected remaining to be 0, got %d", remaining)
		}
		allowed, remaining = call(1)
		if allowed {
			t.Fatalf("expected second request to be denied on iteration %d", i)
		}
		if remaining != 0 {
			t.Fatalf("expected remaining to be 0, got %d", remaining)
		}
	}
}
func TestRedisTokenBucketCases(t *testing.T) {
	rdb := redis.NewClient(&redis.Options{
		Addr: "localhost:6379",
	})
	TestBaseTimeDelta := (TestBaseTimeSec * time.Second) + (TestBaseTimeUs * time.Microsecond)
	tests := []struct {
		Name string
		SkipInitialSetup bool
		InitialTime      time.Time
		InitialTokens    int
		ArgTokens int
		TimeDelta time.Duration
		Allowed      bool
		TokensLeft   int64
		ReturnedTime time.Time  
	}{
		{
			Name: "new request with tokens left",
			InitialTokens: 1,
			InitialTime:   TestBaseTime,
			ArgTokens: 1,
			Allowed:    true,
			TokensLeft: 0,
		},
		{
			Name: "new request where no values have been previously set",
			SkipInitialSetup: true,
			ArgTokens: 1,
			TimeDelta: TestBaseTimeDelta,
			Allowed:    true,
			TokensLeft: TestCapacity - 1,
		},
		{
			Name: "allows for a new request where values were set a long time ago",
			InitialTokens: 0,
			InitialTime:   time.Time{},
			ArgTokens: 1,
			TimeDelta: TestBaseTimeDelta,
			Allowed:    true,
			TokensLeft: TestCapacity - 1,
		},
		{
			Name: "allows for a new request that goes back in time but leaves existing values unchanged",
			InitialTokens: 1,
			InitialTime:   TestBaseTime,
			ArgTokens: 1,
			TimeDelta: -TestBaseTimeDelta,
			Allowed:    true,
			TokensLeft: 0,
			ReturnedTime: TestBaseTime,
		},
		{
			Name: "disallows for a new request at same moment without tokens left",
			InitialTokens: 0,
			InitialTime:   TestBaseTime,
			ArgTokens: 1,
			Allowed:    false,
			TokensLeft: 0,
		},
		{
			Name: "allows for a new request without tokens left but time to replenish",
			InitialTokens: 0,
			InitialTime:   TestBaseTime,
			ArgTokens: 1,
			TimeDelta: SingleTokenDripTimeUs * time.Microsecond,
			Allowed:    true,
			TokensLeft: 0,
		},
		{
			Name: "similarly scales up if more time than necessary has passed",
			InitialTokens: 0,
			InitialTime:   TestBaseTime,
			ArgTokens: 1,
			TimeDelta: (TestCapacity - 1) * SingleTokenDripTimeUs * time.Microsecond,
			Allowed:    true,
			TokensLeft: (TestCapacity - 1) - 1,
		},
		{
			Name: "maxes out at the bucket's capacity",
			InitialTokens: 0,
			InitialTime:   TestBaseTime,
			ArgTokens: 1,
			TimeDelta: (TestCapacity + 1) * SingleTokenDripTimeUs * time.Microsecond,
			Allowed:    true,
			TokensLeft: TestCapacity - 1,
		},
		{
			Name: "allows and leaves values unchanged when requesting 0 tokens",
			InitialTokens: 1,
			InitialTime:   TestBaseTime,
			ArgTokens: 0,
			Allowed:    true,
			TokensLeft: 1,
		},
		{
			Name: "allows and leaves values unchanged when requesting 0 tokens at 0 remaining",
			InitialTokens: 0,
			InitialTime:   TestBaseTime,
			ArgTokens: 0,
			Allowed:    true,
			TokensLeft: 0,
		},
		{
			Name: "denies and returns remaining tokens when requesting more than one",
			InitialTokens: 1,
			InitialTime:   TestBaseTime,
			ArgTokens: 2,
			Allowed:    false,
			TokensLeft: 1,
		},
		{
			Name: "works with real timestamps",
			InitialTokens: 0,
			InitialTime:   time.Now(),
			ArgTokens: 1,
			TimeDelta: SingleTokenDripTimeUs * time.Microsecond,
			Allowed:    true,
			TokensLeft: 0,
		},
	}
	for _, test := range tests {
		t.Run(test.Name, func(t *testing.T) {
			keys := makeKeys()
			if test.InitialTime.IsZero() {
				test.InitialTime = time.Unix(0, 0)
			}
			if !test.SkipInitialSetup {
				mustInitKeys(t, rdb, keys, int64(test.InitialTokens), test.InitialTime)
			}
			argTimeSec, argTimeUs := timeToRedisPair(test.InitialTime.Add(test.TimeDelta))
			resp := updateTokenBucketScript.Eval(
				rdb,
				keys.AsSlice(),
				TestReplenishRate,
				TestCapacity,
				test.ArgTokens,
				argTimeSec,
				argTimeUs,
				"true",
			)
			var expectedReturnedTime, expectedReturnedTimeUs int64
			if test.ReturnedTime.IsZero() {
				expectedReturnedTime, expectedReturnedTimeUs = timeToRedisPair(
					test.InitialTime.Add(test.TimeDelta),
				)
			} else {
				expectedReturnedTime, expectedReturnedTimeUs = timeToRedisPair(test.ReturnedTime)
			}
			assertRedisResults(t, resp,
				test.Allowed,
				test.TokensLeft,
				expectedReturnedTime,
				expectedReturnedTimeUs,
			)
		})
	}
}
func TestReturnedTimeMicroseconds(t *testing.T) {
	rdb := redis.NewClient(&redis.Options{
		Addr: "localhost:6379",
	})
	refreshedTime := time.Unix(TestBaseTimeSec, 0)
	keys := makeKeys()
	mustInitKeys(t, rdb, keys, 0, refreshedTime)
	const extraTime = 100
	argTime := refreshedTime.Add((SingleTokenDripTimeUs + extraTime) * time.Microsecond)
	argTimeSec, argTimeUs := timeToRedisPair(argTime)
	resp := updateTokenBucketScript.Eval(
		rdb,
		keys.AsSlice(),
		TestReplenishRate,
		TestCapacity,
		1,
		argTimeSec,
		argTimeUs,
		"true",
	)
	assertRedisResults(t, resp,
		true,
		0,
		argTimeSec,
		argTimeUs-extraTime,
	)
}
func TestReturnedTimeAcrossSecondBoundary(t *testing.T) {
	rdb := redis.NewClient(&redis.Options{
		Addr: "localhost:6379",
	})
	const refreshedTimeUs = MicrosecondsInSecond - 10
	refreshedTime := time.Unix(TestBaseTimeSec, refreshedTimeUs*int64(time.Microsecond/time.Nanosecond))
	keys := makeKeys()
	mustInitKeys(t, rdb, keys, 0, refreshedTime)
	const extraTime = 100
	argTime := time.Unix(TestBaseTimeSec, 0)
	argTime = argTime.Add((refreshedTimeUs + extraTime) * time.Microsecond)
	argTimeSec, argTimeUs := timeToRedisPair(argTime)
	if refreshedTime.Unix() != argTime.Unix()-1 {
		t.Errorf("expected refreshed time to be 1 second before the argument time: %d != %d - 1",
			refreshedTime.Unix(),
			argTime.Unix()-1,
		)
	}
	resp := updateTokenBucketScript.Eval(
		rdb,
		keys.AsSlice(),
		TestReplenishRate,
		TestCapacity,
		1,
		argTimeSec,
		argTimeUs,
		"true",
	)
	assertRedisResults(t, resp,
		false,
		0,
		argTimeSec-1,
		MicrosecondsInSecond+(argTimeUs-extraTime),
	)
	argTimeUs += SingleTokenDripTimeUs - extraTime
	resp = updateTokenBucketScript.Eval(
		rdb,
		keys.AsSlice(),
		TestReplenishRate,
		TestCapacity,
		1,
		argTimeSec,
		argTimeUs,
		"true",
	)
	assertRedisResults(t, resp,
		true,
		0,
		argTimeSec,
		argTimeUs,
	)
}
func TestHighReplenishRate(t *testing.T) {
	rdb := redis.NewClient(&redis.Options{
		Addr: "localhost:6379",
	})
	keys := makeKeys()
	mustInitKeys(t, rdb, keys, 0, TestBaseTime)
	const (
		tokensToAdd              = 10
		testReplenishRate        = 5000
		testCapacity             = 5000
		singleTokenReplenishTime = MicrosecondsInSecond / testReplenishRate
	)
	argTime := TestBaseTime.Add(singleTokenReplenishTime * tokensToAdd * time.Microsecond)
	argTimeSec, argTimeUs := timeToRedisPair(argTime)
	resp := updateTokenBucketScript.Eval(
		rdb,
		keys.AsSlice(),
		testReplenishRate,
		testCapacity,
		1,
		argTimeSec,
		argTimeUs,
		"true",
	)
	assertRedisResults(t, resp,
		true,
		tokensToAdd-1,
		argTimeSec,
		argTimeUs,
	)
}