linearvoid/backend/chat/src/ratelimit/ratelimit_test.go

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")
}