tokenlimit_test.go 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788
  1. package limit
  2. import (
  3. "testing"
  4. "time"
  5. "git.i2edu.net/i2/go-zero/core/logx"
  6. "git.i2edu.net/i2/go-zero/core/stores/redis"
  7. "git.i2edu.net/i2/go-zero/core/stores/redis/redistest"
  8. "github.com/alicebob/miniredis/v2"
  9. "github.com/stretchr/testify/assert"
  10. )
  11. func init() {
  12. logx.Disable()
  13. }
  14. func TestTokenLimit_Rescue(t *testing.T) {
  15. s, err := miniredis.Run()
  16. assert.Nil(t, err)
  17. const (
  18. total = 100
  19. rate = 5
  20. burst = 10
  21. )
  22. l := NewTokenLimiter(rate, burst, redis.NewRedis(s.Addr(), redis.NodeType), "tokenlimit")
  23. s.Close()
  24. var allowed int
  25. for i := 0; i < total; i++ {
  26. time.Sleep(time.Second / time.Duration(total))
  27. if i == total>>1 {
  28. assert.Nil(t, s.Restart())
  29. }
  30. if l.Allow() {
  31. allowed++
  32. }
  33. // make sure start monitor more than once doesn't matter
  34. l.startMonitor()
  35. }
  36. assert.True(t, allowed >= burst+rate)
  37. }
  38. func TestTokenLimit_Take(t *testing.T) {
  39. store, clean, err := redistest.CreateRedis()
  40. assert.Nil(t, err)
  41. defer clean()
  42. const (
  43. total = 100
  44. rate = 5
  45. burst = 10
  46. )
  47. l := NewTokenLimiter(rate, burst, store, "tokenlimit")
  48. var allowed int
  49. for i := 0; i < total; i++ {
  50. time.Sleep(time.Second / time.Duration(total))
  51. if l.Allow() {
  52. allowed++
  53. }
  54. }
  55. assert.True(t, allowed >= burst+rate)
  56. }
  57. func TestTokenLimit_TakeBurst(t *testing.T) {
  58. store, clean, err := redistest.CreateRedis()
  59. assert.Nil(t, err)
  60. defer clean()
  61. const (
  62. total = 100
  63. rate = 5
  64. burst = 10
  65. )
  66. l := NewTokenLimiter(rate, burst, store, "tokenlimit")
  67. var allowed int
  68. for i := 0; i < total; i++ {
  69. if l.Allow() {
  70. allowed++
  71. }
  72. }
  73. assert.True(t, allowed >= burst)
  74. }