| 34 | func (r *mockResource) Classes() []string { return r.classes } |
| 35 | |
| 36 | func TestRateLimit(t *testing.T) { |
| 37 | a := assertions.New(t) |
| 38 | |
| 39 | limiter, err := ratelimit.New(test.Context(), config.RateLimiting{ |
| 40 | Profiles: []config.RateLimitingProfile{ |
| 41 | { |
| 42 | Name: "Default profile", |
| 43 | MaxPerMin: maxRate, |
| 44 | MaxBurst: maxRate, |
| 45 | Associations: []string{"default"}, |
| 46 | }, |
| 47 | { |
| 48 | Name: "Multiple associations", |
| 49 | MaxPerMin: maxRate, |
| 50 | MaxBurst: maxRate, |
| 51 | Associations: []string{"assoc1", "assoc2"}, |
| 52 | }, |
| 53 | { |
| 54 | Name: "Override", |
| 55 | MaxPerMin: overrideRate, |
| 56 | MaxBurst: overrideRate, |
| 57 | Associations: []string{"override"}, |
| 58 | }, |
| 59 | }, |
| 60 | }, config.BlobConfig{}, nil) |
| 61 | if !a.So(err, should.BeNil) { |
| 62 | t.FailNow() |
| 63 | } |
| 64 | |
| 65 | t.Run("Limit", func(t *testing.T) { |
| 66 | for _, resource := range []*mockResource{ |
| 67 | {"key1", []string{"default"}}, |
| 68 | {"key2", []string{"default"}}, |
| 69 | } { |
| 70 | for i := range maxRate { |
| 71 | limit, result := limiter.RateLimit(resource) |
| 72 | |
| 73 | a.So(limit, should.BeFalse) |
| 74 | a.So(result.Limit, should.Equal, maxRate) |
| 75 | a.So(result.Remaining, should.Equal, maxRate-i-1) |
| 76 | } |
| 77 | |
| 78 | limit, result := limiter.RateLimit(resource) |
| 79 | a.So(limit, should.BeTrue) |
| 80 | a.So(result.Limit, should.Equal, maxRate) |
| 81 | a.So(result.Remaining, should.Equal, 0) |
| 82 | a.So(result.ResetAfter, should.NotBeZeroValue) |
| 83 | a.So(result.RetryAfter, should.NotBeZeroValue) |
| 84 | } |
| 85 | }) |
| 86 | |
| 87 | t.Run("Profile", func(t *testing.T) { |
| 88 | for _, tc := range []struct { |
| 89 | name string |
| 90 | classes []string |
| 91 | validate func(t *testing.T, limit bool, result ratelimit.Result) |
| 92 | }{ |
| 93 | { |