@@ -3,7 +3,6 @@ package main
33import (
44 "context"
55 "net/netip"
6- "sync"
76 "testing"
87 "time"
98
@@ -14,93 +13,67 @@ import (
1413func TestIPThrottler_RegisterNewRequest (t * testing.T ) {
1514 t .Parallel ()
1615
17- t .Run ("valid number of requests " , func (t * testing.T ) {
16+ t .Run ("first request allowed " , func (t * testing.T ) {
1817 t .Parallel ()
1918
20- addr , err := netip .ParseAddr ("127.0.0.1" )
21- require .NoError (t , err )
19+ addr := netip .MustParseAddr ("127.0.0.1" )
2220
23- // Create the IP throttler
2421 th := newIPThrottler (defaultRateLimitInterval , defaultCleanTimeout )
2522
26- // Register < max requests
27- for i := uint64 (0 ); i < maxRequestsPerMinute ; i ++ {
28- assert .NoError (t , th .registerNewRequest (addr ))
29- }
23+ assert .NoError (t , th .registerNewRequest (addr ))
3024 })
3125
32- t .Run ("exceeded number of requests " , func (t * testing.T ) {
26+ t .Run ("second request rejected " , func (t * testing.T ) {
3327 t .Parallel ()
3428
35- addr , err := netip .ParseAddr ("127.0.0.1" )
36- require .NoError (t , err )
29+ addr := netip .MustParseAddr ("127.0.0.1" )
3730
38- // Create the IP throttler
3931 th := newIPThrottler (defaultRateLimitInterval , defaultCleanTimeout )
4032
41- // Register max requests
42- for i := uint64 (0 ); i < maxRequestsPerMinute ; i ++ {
43- assert .NoError (t , th .registerNewRequest (addr ))
44- }
33+ require .NoError (t , th .registerNewRequest (addr ))
4534
46- // Attempt to register an additional request
4735 assert .ErrorIs (t , th .registerNewRequest (addr ), errInvalidNumberOfRequests )
4836 })
4937}
5038
51- func TestIPThrottler_RequestsThrottled (t * testing.T ) {
39+ func TestIPThrottler_SecondRequestRejected (t * testing.T ) {
5240 t .Parallel ()
5341
54- var (
55- cleanupInterval = time .Millisecond * 100
42+ addr := netip .MustParseAddr ("192.168.1.1" )
5643
57- requestInterval = 3 * cleanupInterval // requests triggered after ~5 cleans
58- numRequests = maxRequestsPerMinute * 2 // number of request loops
59- )
44+ // Use a long interval so no tokens regenerate during the test
45+ th := newIPThrottler (time .Hour , defaultCleanTimeout )
6046
61- addr , err := netip . ParseAddr ( "127.0.0.1" )
62- require .NoError (t , err )
47+ // First request must succeed
48+ require .NoError (t , th . registerNewRequest ( addr ) )
6349
64- // Create the IP throttler
65- th := newIPThrottler (defaultRateLimitInterval , cleanupInterval )
66-
67- ctx , cancelFn := context .WithCancel (context .Background ())
68- defer cancelFn ()
50+ // Second request from the same IP must be rejected
51+ assert .ErrorIs (t , th .registerNewRequest (addr ), errInvalidNumberOfRequests )
52+ }
6953
70- // Start the throttler (async)
71- th . start ( ctx )
54+ func TestIPThrottler_CleanupAllowsNewRequest ( t * testing. T ) {
55+ t . Parallel ( )
7256
73- var wg sync. WaitGroup
57+ cleanupInterval := time . Millisecond * 100
7458
75- wg . Add ( 1 )
59+ addr := netip . MustParseAddr ( "127.0.0.1" )
7660
77- go func () {
78- defer wg .Done ()
61+ // Rate interval is long so tokens won't regenerate on their own;
62+ // only cleanup (removing the stale entry) should allow a new request.
63+ th := newIPThrottler (time .Hour , cleanupInterval )
7964
80- var (
81- requestsSent = 0
82- ticker = time .NewTicker (requestInterval )
83- )
65+ ctx , cancelFn := context .WithCancel (context .Background ())
66+ defer cancelFn ()
8467
85- for {
86- select {
87- case <- ctx .Done ():
88- return
89- case <- ticker .C :
90- // Fill out the request count for the address
91- for i := uint64 (0 ); i < maxRequestsPerMinute ; i ++ {
92- require .NoError (t , th .registerNewRequest (addr ))
93- }
68+ th .start (ctx )
9469
95- requestsSent += maxRequestsPerMinute
70+ // First request succeeds, second is rejected
71+ require .NoError (t , th .registerNewRequest (addr ))
72+ require .ErrorIs (t , th .registerNewRequest (addr ), errInvalidNumberOfRequests )
9673
97- if requestsSent == numRequests {
98- // Loops done
99- return
100- }
101- }
102- }
103- }()
74+ // Wait for the cleanup cycle to evict the stale entry
75+ time .Sleep (cleanupInterval * 3 )
10476
105- wg .Wait ()
77+ // After cleanup the IP entry is gone, so a new request succeeds
78+ assert .NoError (t , th .registerNewRequest (addr ))
10679}
0 commit comments