| 1 | // SPDX-License-Identifier: AGPL-3.0-or-later |
| 2 | |
| 3 | package throttle |
| 4 | |
| 5 | import ( |
| 6 | "context" |
| 7 | "testing" |
| 8 | "time" |
| 9 | |
| 10 | "github.com/tenseleyFlow/shithub/internal/testing/dbtest" |
| 11 | ) |
| 12 | |
| 13 | func TestLimiter_HitAndThrottle(t *testing.T) { |
| 14 | t.Parallel() |
| 15 | pool := dbtest.NewTestDB(t) |
| 16 | ctx := context.Background() |
| 17 | l := NewLimiter() |
| 18 | |
| 19 | p := Limit{Scope: "login", Identifier: "ip:1.2.3.4|alice", Max: 3, Window: time.Hour} |
| 20 | |
| 21 | for i := 1; i <= 3; i++ { |
| 22 | if err := l.Hit(ctx, pool, p); err != nil { |
| 23 | t.Fatalf("hit %d: %v", i, err) |
| 24 | } |
| 25 | } |
| 26 | err := l.Hit(ctx, pool, p) |
| 27 | if !IsThrottled(err) { |
| 28 | t.Fatalf("4th hit: expected throttled, got %v", err) |
| 29 | } |
| 30 | } |
| 31 | |
| 32 | func TestLimiter_Reset(t *testing.T) { |
| 33 | t.Parallel() |
| 34 | pool := dbtest.NewTestDB(t) |
| 35 | ctx := context.Background() |
| 36 | l := NewLimiter() |
| 37 | |
| 38 | p := Limit{Scope: "login", Identifier: "ip:1.2.3.4|bob", Max: 1, Window: time.Hour} |
| 39 | |
| 40 | if err := l.Hit(ctx, pool, p); err != nil { |
| 41 | t.Fatalf("first hit: %v", err) |
| 42 | } |
| 43 | if err := l.Hit(ctx, pool, p); !IsThrottled(err) { |
| 44 | t.Fatalf("second hit before reset: expected throttled, got %v", err) |
| 45 | } |
| 46 | if err := l.Reset(ctx, pool, p.Scope, p.Identifier); err != nil { |
| 47 | t.Fatalf("reset: %v", err) |
| 48 | } |
| 49 | if err := l.Hit(ctx, pool, p); err != nil { |
| 50 | t.Fatalf("hit after reset: %v", err) |
| 51 | } |
| 52 | } |
| 53 | |
| 54 | func TestLimiter_WindowReset(t *testing.T) { |
| 55 | t.Parallel() |
| 56 | pool := dbtest.NewTestDB(t) |
| 57 | ctx := context.Background() |
| 58 | l := NewLimiter() |
| 59 | |
| 60 | // Window is short enough that the second hit lands in a brand-new |
| 61 | // window. The bump query resets the counter when the existing window |
| 62 | // started before (now - Window). Use a generous sleep so clock |
| 63 | // granularity / connection latency between the Go cutoff and the PG |
| 64 | // now() can't make the comparison ambiguous. |
| 65 | p := Limit{Scope: "login", Identifier: "ip:1.2.3.4|carol", Max: 1, Window: 200 * time.Millisecond} |
| 66 | |
| 67 | if err := l.Hit(ctx, pool, p); err != nil { |
| 68 | t.Fatalf("first hit: %v", err) |
| 69 | } |
| 70 | time.Sleep(500 * time.Millisecond) |
| 71 | if err := l.Hit(ctx, pool, p); err != nil { |
| 72 | t.Fatalf("hit after window: expected fresh window, got %v", err) |
| 73 | } |
| 74 | } |
| 75 |