go-omni/middleware/redis/store_test.go

package redis

import (
	"context"
	"fmt"
	omni "git.yuntuo.hk/k2safe/go-omni"
	"git.yuntuo.hk/k2safe/go-omni/httpx"
	"git.yuntuo.hk/k2safe/go-omni/middleware"
	goredis "github.com/redis/go-redis/v9"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func TestSharedRedisIntegration(t *testing.T) {
	address := os.Getenv("GO_OMNI_REDIS_TEST_ADDR")
	if address == "" {
		t.Skip("requires isolated Redis")
	}
	client := goredis.NewClient(&goredis.Options{Addr: address, PoolSize: 8, MaxRetries: -1, ContextTimeoutEnabled: true, ReadTimeout: time.Second, WriteTimeout: time.Second})
	defer client.Close()
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	prefix := fmt.Sprint("shared:", time.Now().UnixNano())
	s, err := NewIdempotencyStore(client, prefix, IdempotencyConfig{MaxEntries: 256})
	if err != nil {
		t.Fatal(err)
	}
	defer client.Del(context.Background(), s.keys...)
	key, fp := strings.Repeat("a", 64), strings.Repeat("b", 64)
	var claims atomic.Int32
	var wg sync.WaitGroup
	owners := make(chan string, 32)
	for i := 0; i < 32; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			owner := fmt.Sprintf("%032x", i)
			_, ok, err := s.Claim(ctx, key, fp, owner)
			if err != nil {
				t.Error(err)
			}
			if ok {
				claims.Add(1)
				owners <- owner
			}
		}(i)
	}
	wg.Wait()
	if claims.Load() != 1 {
		t.Fatal(claims.Load())
	}
	owner := <-owners
	if err = s.Complete(ctx, key, "wrong", nil, time.Minute); err == nil {
		t.Fatal("foreign completion accepted")
	}
	if err = s.Complete(ctx, key, owner, &middleware.StoredResponse{Status: 201, Headers: http.Header{}, Body: []byte{0, 1, 255}}, time.Minute); err != nil {
		t.Fatal(err)
	}
	record, ok, err := s.Claim(ctx, key, fp, strings.Repeat("c", 32))
	if err != nil || ok || record.State != "completed" || len(record.Response.Body) != 3 {
		t.Fatal(record, ok, err)
	}
	// Two independently constructed apps use the same backend, so only one runs.
	var executions atomic.Int32
	makeApp := func() *omni.App {
		backend, _ := NewIdempotencyStore(client, prefix, IdempotencyConfig{MaxEntries: 256})
		idem, err := middleware.NewIdempotency(middleware.IdempotencyConfig{Store: backend, Scope: func(*httpx.Context) (string, error) { return "verifiedTenant", nil }})
		if err != nil {
			t.Fatal(err)
		}
		a := omni.New(omni.WithLogger(false))
		a.POST("/submit", func(c *httpx.Context) error { executions.Add(1); return c.OK("saved") }, idem.Middleware())
		return a
	}
	a, b := makeApp(), makeApp()
	request := func(a *omni.App, body string) *httptest.ResponseRecorder {
		r := httptest.NewRequest("POST", "/submit", strings.NewReader(body))
		r.Header.Set("Idempotency-Key", "stableKey")
		w := httptest.NewRecorder()
		a.ServeHTTP(w, r)
		return w
	}
	first, second := request(a, `{}`), request(b, `{}`)
	if first.Code != 200 || second.Code != 200 || first.Body.String() != second.Body.String() || second.Header().Get("Idempotency-Replayed") != "true" || executions.Load() != 1 {
		t.Fatal(first.Code, second.Code, second.Body.String(), executions.Load())
	}
	if request(b, `{"changed":true}`).Code != 409 {
		t.Fatal("conflict accepted")
	}
	// Shared limits do not multiply when a second replica is added.
	limitPrefix := prefix + "limit"
	cfg := middleware.RateLimitConfig{Rate: .01, Burst: 2, Key: func(*httpx.Context) string { return "tenant" }}
	r1, _ := NewRateLimiter(client, limitPrefix, cfg)
	r2, _ := NewRateLimiter(client, limitPrefix, cfg)
	base, _ := namespace(limitPrefix, "rate")
	defer client.Del(context.Background(), base+":buckets", base+":seen", base+":config")
	limited := func(l *middleware.RateLimiter) int {
		a := omni.New(omni.WithLogger(false))
		a.GET("/", func(c *httpx.Context) error { return c.OK("ok") }, l.Middleware())
		w := httptest.NewRecorder()
		a.ServeHTTP(w, httptest.NewRequest("GET", "/", nil))
		return w.Code
	}
	if limited(r1) != 200 || limited(r2) != 200 || limited(r1) != 429 {
		t.Fatal("limits multiplied across instances")
	}
	// Redis failure must fail closed before executing a new request.
	client.Close()
	if status := limited(r2); status != 503 {
		t.Fatal(status)
	}
}

func TestExpiredTargetBeyondCleanupBatch(t *testing.T) {
	address := os.Getenv("GO_OMNI_REDIS_TEST_ADDR")
	if address == "" {
		t.Skip("requires isolated Redis")
	}
	client := goredis.NewClient(&goredis.Options{Addr: address})
	defer client.Close()
	s, _ := NewIdempotencyStore(client, fmt.Sprint("expiry:", time.Now().UnixNano()), IdempotencyConfig{MaxEntries: 256})
	defer client.Del(context.Background(), s.keys...)
	ctx := context.Background()
	owner := strings.Repeat("1", 32)
	fp := strings.Repeat("2", 64)
	for i := 0; i < 140; i++ {
		key := fmt.Sprintf("%064x", i)
		if _, ok, err := s.Claim(ctx, key, fp, owner); err != nil || !ok {
			t.Fatal(err)
		}
		if err := s.Complete(ctx, key, owner, nil, time.Minute); err != nil {
			t.Fatal(err)
		}
		client.ZAdd(ctx, s.keys[1], goredis.Z{Score: 0, Member: key})
	}
	if _, ok, err := s.Claim(ctx, fmt.Sprintf("%064x", 139), fp, owner); err != nil || !ok {
		t.Fatal("expired target retained beyond TTL", ok, err)
	}
}