121 lines
3.0 KiB
Go
121 lines
3.0 KiB
Go
package stocklock
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
"fmt"
|
||
"os"
|
||
"sync"
|
||
"sync/atomic"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/redis/go-redis/v9"
|
||
)
|
||
|
||
// 需要 Redis：REDIS_URL=redis://localhost:6379/0 go test ./...
|
||
func client(t *testing.T) (*redis.Client, string) {
|
||
t.Helper()
|
||
url := os.Getenv("REDIS_URL")
|
||
if url == "" {
|
||
t.Skip("未设置 REDIS_URL，跳过集成测试")
|
||
}
|
||
opt, err := redis.ParseURL(url)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
rdb := redis.NewClient(opt)
|
||
prefix := fmt.Sprintf("test:%d:", time.Now().UnixNano())
|
||
t.Cleanup(func() {
|
||
ctx := context.Background()
|
||
iter := rdb.Scan(ctx, 0, prefix+"*", 100).Iterator()
|
||
for iter.Next(ctx) {
|
||
rdb.Del(ctx, iter.Val())
|
||
}
|
||
rdb.Close()
|
||
})
|
||
return rdb, prefix
|
||
}
|
||
|
||
func TestLock(t *testing.T) {
|
||
rdb, p := client(t)
|
||
ctx := context.Background()
|
||
l1, err := TryLock(ctx, rdb, p+"lock", 200*time.Millisecond)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if _, err := TryLock(ctx, rdb, p+"lock", time.Second); !errors.Is(err, ErrLockNotAcquired) {
|
||
t.Fatalf("锁被重复获取：%v", err)
|
||
}
|
||
time.Sleep(300 * time.Millisecond) // 锁过期，别人拿到锁
|
||
l2, err := TryLock(ctx, rdb, p+"lock", time.Second)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if ok, _ := l1.Unlock(ctx); ok {
|
||
t.Fatal("过期的旧持有者不应该删掉新持有者的锁")
|
||
}
|
||
if ok, _ := l2.Unlock(ctx); !ok {
|
||
t.Fatal("持有者应该能释放锁")
|
||
}
|
||
}
|
||
|
||
func TestSlidingWindow(t *testing.T) {
|
||
rdb, p := client(t)
|
||
ctx := context.Background()
|
||
sw := NewSlidingWindow(rdb, 3, time.Second)
|
||
now := time.Now()
|
||
var got []bool
|
||
for i := range 5 {
|
||
ok, err := sw.Allow(ctx, p+"rl", now.Add(time.Duration(i)*10*time.Millisecond))
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
got = append(got, ok)
|
||
}
|
||
if fmt.Sprint(got) != "[true true true false false]" {
|
||
t.Fatalf("窗口内限 3 次，得到 %v", got)
|
||
}
|
||
// 1.1 秒后前面的请求滑出窗口
|
||
if ok, _ := sw.Allow(ctx, p+"rl", now.Add(1100*time.Millisecond)); !ok {
|
||
t.Fatal("窗口滑过后应该放行")
|
||
}
|
||
}
|
||
|
||
func TestReserveNoOversell(t *testing.T) {
|
||
rdb, p := client(t)
|
||
ctx := context.Background()
|
||
sku := p + "sku1"
|
||
rdb.Set(ctx, "stock:"+sku, 10, 0)
|
||
t.Cleanup(func() { rdb.Del(ctx, "stock:"+sku, "stock:"+sku+":buyers") })
|
||
|
||
var success, soldOut, dup atomic.Int32
|
||
var wg sync.WaitGroup
|
||
for i := range 100 { // 100 个用户抢 10 件
|
||
for range 2 { // 每个用户重复点两次
|
||
wg.Add(1)
|
||
go func() {
|
||
defer wg.Done()
|
||
_, err := Reserve(ctx, rdb, sku, fmt.Sprint("u", i))
|
||
switch {
|
||
case err == nil:
|
||
success.Add(1)
|
||
case errors.Is(err, ErrSoldOut):
|
||
soldOut.Add(1)
|
||
case errors.Is(err, ErrDuplicate):
|
||
dup.Add(1)
|
||
default:
|
||
t.Error(err)
|
||
}
|
||
}()
|
||
}
|
||
}
|
||
wg.Wait()
|
||
left, _ := rdb.Get(ctx, "stock:"+sku).Int()
|
||
buyers, _ := rdb.SCard(ctx, "stock:"+sku+":buyers").Result()
|
||
if success.Load() != 10 || left != 0 || buyers != 10 {
|
||
t.Fatalf("成功 %d，剩余 %d，购买人数 %d", success.Load(), left, buyers)
|
||
}
|
||
t.Logf("成功 %d，售罄 %d，重复 %d", success.Load(), soldOut.Load(), dup.Load())
|
||
}
|