86 lines
1.8 KiB
Go
86 lines
1.8 KiB
Go
package trader
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestRateLimiter(t *testing.T) {
|
|
limiter := NewRateLimiter(3, 2)
|
|
|
|
// Should allow first trade
|
|
if !limiter.CanTrade() {
|
|
t.Error("Should allow first trade")
|
|
}
|
|
|
|
limiter.RecordTrade()
|
|
|
|
// Should allow second trade (within parallel limit)
|
|
if !limiter.CanTrade() {
|
|
t.Error("Should allow second trade")
|
|
}
|
|
|
|
limiter.RecordTrade()
|
|
|
|
// Should NOT allow third trade (parallel limit reached)
|
|
if limiter.CanTrade() {
|
|
t.Error("Should NOT allow third trade (parallel limit)")
|
|
}
|
|
|
|
// Release one trade
|
|
limiter.ReleaseTrade()
|
|
|
|
// Should allow trade again
|
|
if !limiter.CanTrade() {
|
|
t.Error("Should allow trade after release")
|
|
}
|
|
|
|
limiter.RecordTrade()
|
|
|
|
// Now at hourly limit (3 trades total, 4th trade)
|
|
limiter.RecordTrade()
|
|
|
|
// Release one trade to free up parallel slot
|
|
limiter.ReleaseTrade()
|
|
|
|
// Should NOT allow (hourly limit reached: 4 >= 3)
|
|
if limiter.CanTrade() {
|
|
t.Error("Should NOT allow trade (hourly limit)")
|
|
}
|
|
|
|
hourly, active := limiter.GetStats()
|
|
if hourly != 4 {
|
|
t.Errorf("Expected 4 hourly trades, got %d", hourly)
|
|
}
|
|
if active != 2 {
|
|
t.Errorf("Expected 2 active trades after one release, got %d", active)
|
|
}
|
|
}
|
|
|
|
func TestRateLimiterHourBoundary(t *testing.T) {
|
|
limiter := NewRateLimiter(3, 5)
|
|
|
|
// Record 3 trades
|
|
for i := 0; i < 3; i++ {
|
|
limiter.RecordTrade()
|
|
}
|
|
|
|
// Should be at limit
|
|
if limiter.CanTrade() {
|
|
t.Error("Should be at hourly limit")
|
|
}
|
|
|
|
// Manually manipulate the hour bucket to simulate time passing
|
|
// In real usage, the hour buckets are cleaned up automatically
|
|
limiter.mu.Lock()
|
|
currentHour := time.Now().Unix() / 3600
|
|
// Clear current hour to simulate new hour
|
|
delete(limiter.hourlyTrades, currentHour)
|
|
limiter.mu.Unlock()
|
|
|
|
// Should allow trades again in new hour
|
|
if !limiter.CanTrade() {
|
|
t.Error("Should allow trades in new hour")
|
|
}
|
|
}
|