traefik-llmgateway/redis_store_test.go

package traefikllmgateway

import (
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestRedisStore_IncrBy_PipelinesIncrbyAndExpire is the brief's Step-1
// pipeline case: incrBy sends INCRBY and EXPIRE as one pipeline and
// returns INCRBY's reply value.
func TestRedisStore_IncrBy_PipelinesIncrbyAndExpire(t *testing.T) {
	ln := newFakeListener(t)
	runFakeRESPServer(t, ln, []respStep{
		{wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")},
		{wantArgs: []string{"INCRBY", "llmgw:user:a:req:day:20260820", "1"}, reply: []byte(":5\r\n")},
		{wantArgs: []string{"EXPIRE", "llmgw:user:a:req:day:20260820", "60"}, reply: []byte(":1\r\n")},
	})

	store := newRedisStore(newRESPClient(ln.Addr().String(), "", 0))
	v, err := store.incrBy("llmgw:user:a:req:day:20260820", 1, 60*time.Second)
	if err != nil {
		t.Fatalf("incrBy: %v", err)
	}
	if v != 5 {
		t.Errorf("incrBy = %d, want 5", v)
	}
}

// TestRedisStore_IncrMulti_PipelinesAllPairsInOneCall is the perf-review
// (2026-08-21) case: incrMulti sends every entry's INCRBY+EXPIRE pair as
// ONE pipeline — no SELECT (or any other command) interleaved between
// entries, which would only happen if incrMulti dialled a fresh
// connection or made a separate call per entry instead of the single
// s.client.pipeline call it actually makes — and returns each entry's
// INCRBY reply in the same order as entries, not the reply's position in
// the full 2*N command/reply stream.
func TestRedisStore_IncrMulti_PipelinesAllPairsInOneCall(t *testing.T) {
	ln := newFakeListener(t)
	runFakeRESPServer(t, ln, []respStep{
		{wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")},
		{wantArgs: []string{"INCRBY", "llmgw:user:a:req:hour:2026082014", "1"}, reply: []byte(":11\r\n")},
		{wantArgs: []string{"EXPIRE", "llmgw:user:a:req:hour:2026082014", "172800"}, reply: []byte(":1\r\n")},
		{wantArgs: []string{"INCRBY", "llmgw:user:a:tokin:day:20260820", "40"}, reply: []byte(":140\r\n")},
		{wantArgs: []string{"EXPIRE", "llmgw:user:a:tokin:day:20260820", "3024000"}, reply: []byte(":1\r\n")},
		{wantArgs: []string{"INCRBY", "llmgw:total:all:cost:month:202608", "500"}, reply: []byte(":9500\r\n")},
		{wantArgs: []string{"EXPIRE", "llmgw:total:all:cost:month:202608", "34560000"}, reply: []byte(":1\r\n")},
	})

	store := newRedisStore(newRESPClient(ln.Addr().String(), "", 0))
	entries := []counterIncr{
		{key: "llmgw:user:a:req:hour:2026082014", delta: 1, ttl: hourWindowTTL},
		{key: "llmgw:user:a:tokin:day:20260820", delta: 40, ttl: dayWindowTTL},
		{key: "llmgw:total:all:cost:month:202608", delta: 500, ttl: monthWindowTTL},
	}
	got, err := store.incrMulti(entries)
	if err != nil {
		t.Fatalf("incrMulti: %v", err)
	}
	want := []int64{11, 140, 9500}
	if len(got) != len(want) {
		t.Fatalf("len(got) = %d, want %d", len(got), len(want))
	}
	for i, w := range want {
		if got[i] != w {
			t.Errorf("got[%d] = %d, want %d (entry order, not reply-stream position)", i, got[i], w)
		}
	}
}

// TestRedisStore_IncrAndGetMulti_PipelinesEverythingInOneCall is the
// perf-review round-3 (2026-08-22) case: incrAndGetMulti sends every
// entry's INCRBY+EXPIRE pair AND every read key's GET as ONE pipeline —
// no SELECT (or any other command) interleaved, which would only happen
// if it dialled a fresh connection or split the call into more than one
// round trip — and returns entries' INCRBY replies and reads' GET replies
// each in their own given order, not the reply stream's flat position.
func TestRedisStore_IncrAndGetMulti_PipelinesEverythingInOneCall(t *testing.T) {
	ln := newFakeListener(t)
	runFakeRESPServer(t, ln, []respStep{
		{wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")},
		{wantArgs: []string{"INCRBY", "llmgw:user:a:req:min:202608220900", "1"}, reply: []byte(":3\r\n")},
		{wantArgs: []string{"EXPIRE", "llmgw:user:a:req:min:202608220900", "120"}, reply: []byte(":1\r\n")},
		{wantArgs: []string{"INCRBY", "llmgw:user:a:req:day:20260822", "1"}, reply: []byte(":7\r\n")},
		{wantArgs: []string{"EXPIRE", "llmgw:user:a:req:day:20260822", "3024000"}, reply: []byte(":1\r\n")},
		{wantArgs: []string{"GET", "llmgw:user:a:tokin:day:20260822"}, reply: []byte("$3\r\n150\r\n")},
		{wantArgs: []string{"GET", "llmgw:user:a:tokout:day:20260822"}, reply: []byte("$-1\r\n")}, // missing -> 0
		{wantArgs: []string{"GET", "llmgw:user:a:cost:day:20260822"}, reply: []byte("$2\r\n50\r\n")},
	})

	store := newRedisStore(newRESPClient(ln.Addr().String(), "", 0))
	entries := []counterIncr{
		{key: "llmgw:user:a:req:min:202608220900", delta: 1, ttl: minWindowTTL},
		{key: "llmgw:user:a:req:day:20260822", delta: 1, ttl: dayWindowTTL},
	}
	reads := []string{
		"llmgw:user:a:tokin:day:20260822",
		"llmgw:user:a:tokout:day:20260822",
		"llmgw:user:a:cost:day:20260822",
	}
	incrVals, readVals, err := store.incrAndGetMulti(entries, reads)
	if err != nil {
		t.Fatalf("incrAndGetMulti: %v", err)
	}
	assert.Equal(t, []int64{3, 7}, incrVals)
	assert.Equal(t, []int64{150, 0, 50}, readVals)
}

// TestRedisStore_IncrAndGetMulti_EmptyReads asserts a nil/empty reads
// slice (no scope in the request has a token/cost budget configured — the
// common case) still sends and returns the increment half correctly, with
// no GET commands at all.
func TestRedisStore_IncrAndGetMulti_EmptyReads(t *testing.T) {
	ln := newFakeListener(t)
	runFakeRESPServer(t, ln, []respStep{
		{wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")},
		{wantArgs: []string{"INCRBY", "k", "1"}, reply: []byte(":1\r\n")},
		{wantArgs: []string{"EXPIRE", "k", "60"}, reply: []byte(":1\r\n")},
	})

	store := newRedisStore(newRESPClient(ln.Addr().String(), "", 0))
	incrVals, readVals, err := store.incrAndGetMulti([]counterIncr{{key: "k", delta: 1, ttl: time.Minute}}, nil)
	if err != nil {
		t.Fatalf("incrAndGetMulti: %v", err)
	}
	if len(incrVals) != 1 || incrVals[0] != 1 {
		t.Errorf("incrVals = %v, want [1]", incrVals)
	}
	if len(readVals) != 0 {
		t.Errorf("readVals = %v, want empty", readVals)
	}
}

// TestRedisStore_IncrBy_RoundsSubSecondTTLUp asserts a ttl under one
// second is never sent to EXPIRE as 0 (which would delete the key
// immediately) — it is rounded up to a 1s floor.
func TestRedisStore_IncrBy_RoundsSubSecondTTLUp(t *testing.T) {
	ln := newFakeListener(t)
	runFakeRESPServer(t, ln, []respStep{
		{wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")},
		{wantArgs: []string{"INCRBY", "k", "1"}, reply: []byte(":1\r\n")},
		{wantArgs: []string{"EXPIRE", "k", "1"}, reply: []byte(":1\r\n")},
	})

	store := newRedisStore(newRESPClient(ln.Addr().String(), "", 0))
	if _, err := store.incrBy("k", 1, 10*time.Millisecond); err != nil {
		t.Fatalf("incrBy: %v", err)
	}
}

// TestRedisStore_Get_MissingKeyReturnsZero is the brief's Step-1 GET case:
// a null bulk reply ("$-1") decodes to 0 with no error.
func TestRedisStore_Get_MissingKeyReturnsZero(t *testing.T) {
	ln := newFakeListener(t)
	runFakeRESPServer(t, ln, []respStep{
		{wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")},
		{wantArgs: []string{"GET", "llmgw:user:a:req:day:20260820"}, reply: []byte("$-1\r\n")},
	})

	store := newRedisStore(newRESPClient(ln.Addr().String(), "", 0))
	v, err := store.get("llmgw:user:a:req:day:20260820")
	if err != nil {
		t.Fatalf("get: %v", err)
	}
	if v != 0 {
		t.Errorf("get = %d, want 0", v)
	}
}

// TestRedisStore_Get_ParsesIntegerValue asserts a present key's bulk-string
// value is parsed as a base-10 integer.
func TestRedisStore_Get_ParsesIntegerValue(t *testing.T) {
	ln := newFakeListener(t)
	runFakeRESPServer(t, ln, []respStep{
		{wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")},
		{wantArgs: []string{"GET", "k"}, reply: []byte("$3\r\n123\r\n")},
	})

	store := newRedisStore(newRESPClient(ln.Addr().String(), "", 0))
	v, err := store.get("k")
	if err != nil {
		t.Fatalf("get: %v", err)
	}
	if v != 123 {
		t.Errorf("get = %d, want 123", v)
	}
}

// TestLimiter_FailOpenFalse_DeadRedisAddress_ReturnsStoreDownViolation is
// the brief's Step-1 case: a limiter backed by a redisStore pointed at an
// address nothing listens on, with failOpen=false, must refuse the
// request with a storeDown violation rather than silently falling back to
// an in-process counter or panicking.
func TestLimiter_FailOpenFalse_DeadRedisAddress_ReturnsStoreDownViolation(t *testing.T) {
	ln := newFakeListener(t)
	deadAddr := ln.Addr().String()
	if err := ln.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	} // nothing listens at deadAddr from here on

	store := newRedisStore(newRESPClient(deadAddr, "", 0))
	l := newLimiter(store, false)
	scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{RequestsPerMinute: 100}}}

	v := l.checkAndCount(scopes)
	if v == nil {
		t.Fatal("want a violation when the store is unreachable and failOpen is false")
	}
	if !v.storeDown {
		t.Errorf("v.storeDown = false, want true")
	}
	if v.message != "limit store unavailable" {
		t.Errorf("v.message = %q, want %q", v.message, "limit store unavailable")
	}
}

// TestLimiter_FailOpenTrue_DeadRedisAddress_FallsBackAndAllows mirrors the
// fail-closed case above with failOpen=true: the same dead address must
// not block the request — the limiter transparently counts against its
// in-process fallback instead.
func TestLimiter_FailOpenTrue_DeadRedisAddress_FallsBackAndAllows(t *testing.T) {
	ln := newFakeListener(t)
	deadAddr := ln.Addr().String()
	if err := ln.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}

	store := newRedisStore(newRESPClient(deadAddr, "", 0))
	l := newLimiter(store, true)
	scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{RequestsPerMinute: 100}}}

	if v := l.checkAndCount(scopes); v != nil {
		t.Fatalf("want no violation with failOpen=true and a dead store, got %+v", v)
	}

	got, err := l.fallback.get(windowKey("user", "u", metricReq, windowDay, l.now()))
	if err != nil {
		t.Fatalf("fallback.get: %v", err)
	}
	if got != 1 {
		t.Errorf("fallback req:day counter = %d, want 1 (the request must have been counted in the fallback)", got)
	}
}

// TestLimiter_FailOpenTrue_HungRedisServer_FallsBackQuickly is review item
// 1's limiter-level proof (round 1), tightened by the round-3 store-down
// latch: a hung (accepts the connection, never replies) Redis-compatible
// server must not make a failOpen=true request wait out several stacked
// multi-second timeouts (the pre-fix behavior for a checkAndCount call
// that performs several store operations). checkAndCount's two store ops
// here (req:min, req:day) would each independently be bounded by one
// respCallTimeout — already a large improvement over the old
// per-operation-multiplied stall — but the store-down latch (limits.go)
// goes further: the first op's failure latches the store down, so the
// second op skips the network call entirely and the whole call costs
// only ~one respCallTimeout, not two.
func TestLimiter_FailOpenTrue_HungRedisServer_FallsBackQuickly(t *testing.T) {
	ln := newHungListener(t)
	store := newRedisStore(newRESPClient(ln.Addr().String(), "", 0))
	l := newLimiter(store, true)
	scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{RequestsPerMinute: 100}}}

	start := time.Now()
	v := l.checkAndCount(scopes)
	elapsed := time.Since(start)

	if v != nil {
		t.Fatalf("want no violation with failOpen=true and a hung store, got %+v", v)
	}
	if elapsed >= 3*time.Second {
		t.Fatalf("elapsed = %v, want < 3s (one respCallTimeout for the first op, the second latched and skipping the network entirely)", elapsed)
	}
}

// TestRedisStore_Get_DownServer_ReturnsError asserts get surfaces a
// transport failure (dead address) as an error, mirroring incrBy's own
// contract for the same failure — the limiter's storeGet relies on this to
// treat a Redis outage as a store error rather than a false zero reading.
func TestRedisStore_Get_DownServer_ReturnsError(t *testing.T) {
	ln := newFakeListener(t)
	deadAddr := ln.Addr().String()
	require.NoError(t, ln.Close())

	store := newRedisStore(newRESPClient(deadAddr, "", 0))
	_, err := store.get("k")
	require.Error(t, err)
	assert.Contains(t, err.Error(), "redisStore: get")
}