redis-rate-limiter/lock/lock_test.go

package lock

import (
	"context"
	"os"
	"testing"
	"time"

	"github.com/joho/godotenv"
	"github.com/redis/go-redis/v9"
)

func setupClient(t *testing.T) *redis.Client {
	t.Helper()

	// 尝试加载 .env（兼容从项目根目录或子目录运行）
	_ = godotenv.Load("../../.env")
	_ = godotenv.Load("../.env")
	_ = godotenv.Load(".env")

	url := os.Getenv("REDIS_URL")
	if url == "" {
		t.Skip("REDIS_URL not set, skipping integration test")
	}

	opt, err := redis.ParseURL(url)
	if err != nil {
		t.Fatalf("Invalid REDIS_URL: %v", err)
	}

	client := redis.NewClient(opt)
	if err := client.Ping(context.Background()).Err(); err != nil {
		t.Fatalf("Redis connection failed: %v", err)
	}

	t.Cleanup(func() { client.Close() })
	return client
}

func TestAcquireAndRelease(t *testing.T) {
	client := setupClient(t)
	ctx := context.Background()

	l := New(client, "test-lock-acquire", "owner-1", 10*time.Second)

	// 清理可能残留的 key
	_ = l.Release(ctx)

	// 第一次获取应该成功
	ok, err := l.Acquire(ctx)
	if err != nil {
		t.Fatalf("Acquire error: %v", err)
	}
	if !ok {
		t.Fatal("Expected to acquire lock")
	}

	// 释放
	if err := l.Release(ctx); err != nil {
		t.Fatalf("Release error: %v", err)
	}

	// 释放后再次获取应该成功
	ok, err = l.Acquire(ctx)
	if err != nil {
		t.Fatalf("Acquire after release error: %v", err)
	}
	if !ok {
		t.Fatal("Expected to acquire lock after release")
	}

	_ = l.Release(ctx)
}

func TestAcquireConflict(t *testing.T) {
	client := setupClient(t)
	ctx := context.Background()

	l1 := New(client, "test-lock-conflict", "owner-1", 10*time.Second)
	l2 := New(client, "test-lock-conflict", "owner-2", 10*time.Second)

	_ = l1.Release(ctx)
	_ = l2.Release(ctx)

	// owner-1 获取锁
	ok, err := l1.Acquire(ctx)
	if err != nil {
		t.Fatalf("Acquire error: %v", err)
	}
	if !ok {
		t.Fatal("Expected owner-1 to acquire lock")
	}

	// owner-2 获取同一把锁应该失败
	ok, err = l2.Acquire(ctx)
	if err != nil {
		t.Fatalf("Acquire error: %v", err)
	}
	if ok {
		t.Fatal("Expected owner-2 to fail acquiring lock")
	}

	// owner-2 释放别人的锁应该失败
	err = l2.Release(ctx)
	if err != ErrLockNotHeld {
		t.Fatalf("Expected ErrLockNotHeld, got: %v", err)
	}

	// owner-1 正常释放
	if err := l1.Release(ctx); err != nil {
		t.Fatalf("Release error: %v", err)
	}
}

func TestRefresh(t *testing.T) {
	client := setupClient(t)
	ctx := context.Background()

	l := New(client, "test-lock-refresh", "owner-1", 5*time.Second)
	_ = l.Release(ctx)

	ok, err := l.Acquire(ctx)
	if err != nil {
		t.Fatalf("Acquire error: %v", err)
	}
	if !ok {
		t.Fatal("Expected to acquire lock")
	}

	// 续期
	ok, err = l.Refresh(ctx)
	if err != nil {
		t.Fatalf("Refresh error: %v", err)
	}
	if !ok {
		t.Fatal("Expected refresh to succeed")
	}

	_ = l.Release(ctx)
}

func TestAcquireWithRetry(t *testing.T) {
	client := setupClient(t)
	ctx := context.Background()

	l1 := New(client, "test-lock-retry", "owner-1", 10*time.Second)
	l2 := New(client, "test-lock-retry", "owner-2", 10*time.Second)

	_ = l1.Release(ctx)

	// owner-1 先获取锁
	ok, _ := l1.Acquire(ctx)
	if !ok {
		t.Fatal("Expected owner-1 to acquire lock")
	}

	// owner-2 带重试，100ms 后 owner-1 释放，owner-2 应该能拿到
	go func() {
		time.Sleep(100 * time.Millisecond)
		_ = l1.Release(ctx)
	}()

	ok, err := l2.AcquireWithRetry(ctx, 50*time.Millisecond, 10)
	if err != nil {
		t.Fatalf("AcquireWithRetry error: %v", err)
	}
	if !ok {
		t.Fatal("Expected owner-2 to acquire lock after retry")
	}

	_ = l2.Release(ctx)
}