package ratelimit import ( "context" "sync" "testing" "time" ) // mockRedis implements RedisClient for testing. type mockRedis struct { mu sync.Mutex store map[string]int64 } func newMockRedis() *mockRedis { return &mockRedis{store: make(map[string]int64)} } func (m *mockRedis) Incr(ctx context.Context, key string) (int64, error) { m.mu.Lock() defer m.mu.Unlock() m.store[key]++ return m.store[key], nil } func (m *mockRedis) Expire(ctx context.Context, key string, expiration time.Duration) error { return nil } func (m *mockRedis) Decr(ctx context.Context, key string) (int64, error) { m.mu.Lock() defer m.mu.Unlock() m.store[key]-- return m.store[key], nil } func TestCheckMessageRate(t *testing.T) { redis := newMockRedis() ctx := context.Background() // First 10 messages allowed for i := 0; i < 10; i++ { allowed, err := CheckMessageRate(ctx, redis, "user1", 10, 10*time.Second) if err != nil { t.Fatalf("CheckMessageRate #%d error = %v", i+1, err) } if !allowed { t.Errorf("CheckMessageRate #%d = false, want true", i+1) } } // 11th message blocked allowed, err := CheckMessageRate(ctx, redis, "user1", 10, 10*time.Second) if err != nil { t.Fatalf("CheckMessageRate #11 error = %v", err) } if allowed { t.Error("CheckMessageRate #11 = true, want false (rate limited)") } // Different user not affected allowed, err = CheckMessageRate(ctx, redis, "user2", 10, 10*time.Second) if err != nil { t.Fatalf("CheckMessageRate user2 error = %v", err) } if !allowed { t.Error("CheckMessageRate user2 = false, want true") } } func TestCheckMessageRate_NilRedis(t *testing.T) { allowed, err := CheckMessageRate(context.Background(), nil, "user1", 10, 10*time.Second) if err != nil { t.Fatalf("CheckMessageRate nil redis error = %v", err) } if !allowed { t.Error("CheckMessageRate nil redis = false, want true (fail open)") } } func TestCheckLoginRate(t *testing.T) { redis := newMockRedis() ctx := context.Background() // First 5 attempts allowed for i := 0; i < 5; i++ { allowed, err := CheckLoginRate(ctx, redis, "192.168.1.1") if err != nil { t.Fatalf("CheckLoginRate #%d error = %v", i+1, err) } if !allowed { t.Errorf("CheckLoginRate #%d = false, want true", i+1) } } // 6th attempt blocked allowed, err := CheckLoginRate(ctx, redis, "192.168.1.1") if err != nil { t.Fatalf("CheckLoginRate #6 error = %v", err) } if allowed { t.Error("CheckLoginRate #6 = true, want false (rate limited)") } // Different IP not affected allowed, err = CheckLoginRate(ctx, redis, "10.0.0.1") if err != nil { t.Fatalf("CheckLoginRate other IP error = %v", err) } if !allowed { t.Error("CheckLoginRate other IP = false, want true") } } func TestCheckLoginRate_NilRedis(t *testing.T) { allowed, err := CheckLoginRate(context.Background(), nil, "192.168.1.1") if err != nil { t.Fatalf("CheckLoginRate nil redis error = %v", err) } if !allowed { t.Error("CheckLoginRate nil redis = false, want true (fail open)") } } func TestCheckWSConnRate(t *testing.T) { redis := newMockRedis() ctx := context.Background() // First 5 connections allowed for i := 0; i < 5; i++ { allowed, err := CheckWSConnRate(ctx, redis, "192.168.1.1") if err != nil { t.Fatalf("CheckWSConnRate #%d error = %v", i+1, err) } if !allowed { t.Errorf("CheckWSConnRate #%d = false, want true", i+1) } } // 6th connection blocked allowed, err := CheckWSConnRate(ctx, redis, "192.168.1.1") if err != nil { t.Fatalf("CheckWSConnRate #6 error = %v", err) } if allowed { t.Error("CheckWSConnRate #6 = true, want false (rate limited)") } } func TestCheckWSConnRate_NilRedis(t *testing.T) { allowed, err := CheckWSConnRate(context.Background(), nil, "192.168.1.1") if err != nil { t.Fatalf("CheckWSConnRate nil redis error = %v", err) } if !allowed { t.Error("CheckWSConnRate nil redis = false, want true (fail open)") } } func TestDecrWSConn(t *testing.T) { redis := newMockRedis() ctx := context.Background() // Increment to 5 for i := 0; i < 5; i++ { CheckWSConnRate(ctx, redis, "192.168.1.1") } // 6th should be blocked allowed, _ := CheckWSConnRate(ctx, redis, "192.168.1.1") if allowed { t.Error("6th connection should be blocked") } // Decrement twice (simulate 2 disconnects) DecrWSConn(ctx, redis, "192.168.1.1") DecrWSConn(ctx, redis, "192.168.1.1") // Now should be allowed again (counter 6→4, then incr→5) allowed, _ = CheckWSConnRate(ctx, redis, "192.168.1.1") if !allowed { t.Error("after 2x DecrWSConn, connection should be allowed") } } func TestDecrWSConn_NilRedis(t *testing.T) { // Should not panic DecrWSConn(context.Background(), nil, "192.168.1.1") }