CLIProxyAPI/internal/runtime/executor/helps/session_id_cache_test.go

package helps

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/google/uuid"
)

func resetSessionIDCache() {
	sessionIDCacheMu.Lock()
	sessionIDCache = make(map[string]sessionIDCacheEntry)
	sessionIDCacheMu.Unlock()
}

type fakeClaudeIDKVClient struct {
	values        map[string][]byte
	getErr        error
	setErr        error
	expireErr     error
	setNoPersist  bool
	getCount      int
	setCount      int
	expireCount   int
	lastSetTTL    time.Duration
	lastExpireTTL time.Duration
}

func newFakeClaudeIDKVClient() *fakeClaudeIDKVClient {
	return &fakeClaudeIDKVClient{values: make(map[string][]byte)}
}

func (c *fakeClaudeIDKVClient) KVGet(_ context.Context, key string) ([]byte, bool, error) {
	c.getCount++
	if c.getErr != nil {
		return nil, false, c.getErr
	}
	value, ok := c.values[key]
	if !ok {
		return nil, false, nil
	}
	return append([]byte(nil), value...), true, nil
}

func (c *fakeClaudeIDKVClient) KVSetNX(_ context.Context, key string, value []byte, ttl time.Duration) (bool, error) {
	c.setCount++
	c.lastSetTTL = ttl
	if c.setErr != nil {
		return false, c.setErr
	}
	if _, ok := c.values[key]; ok {
		return false, nil
	}
	if !c.setNoPersist {
		c.values[key] = append([]byte(nil), value...)
	}
	return true, nil
}

func (c *fakeClaudeIDKVClient) KVExpire(_ context.Context, _ string, ttl time.Duration) (bool, error) {
	c.expireCount++
	c.lastExpireTTL = ttl
	if c.expireErr != nil {
		return false, c.expireErr
	}
	return true, nil
}

func useFakeClaudeIDKVClient(t *testing.T, client *fakeClaudeIDKVClient, homeMode bool, errClient error) {
	t.Helper()
	previous := currentClaudeIDKVClient
	currentClaudeIDKVClient = func() (claudeIDKVClient, bool, error) {
		return client, homeMode, errClient
	}
	t.Cleanup(func() {
		currentClaudeIDKVClient = previous
	})
}

func TestCachedSessionIDRequiredHomeReusesKVAcrossLocalCacheReset(t *testing.T) {
	resetSessionIDCache()
	client := newFakeClaudeIDKVClient()
	useFakeClaudeIDKVClient(t, client, true, nil)

	first, errFirst := CachedSessionIDRequired(context.Background(), "api-key-1")
	if errFirst != nil {
		t.Fatalf("CachedSessionIDRequired() first error = %v", errFirst)
	}
	resetSessionIDCache()
	second, errSecond := CachedSessionIDRequired(context.Background(), "api-key-1")
	if errSecond != nil {
		t.Fatalf("CachedSessionIDRequired() second error = %v", errSecond)
	}
	if first != second {
		t.Fatalf("session id = %q then %q, want same Home KV value", first, second)
	}
	if _, errParse := uuid.Parse(first); errParse != nil {
		t.Fatalf("session id %q is not a UUID: %v", first, errParse)
	}
	if client.setCount != 1 {
		t.Fatalf("KVSetNX count = %d, want 1", client.setCount)
	}
	if client.expireCount != 1 || client.lastExpireTTL != sessionIDTTL {
		t.Fatalf("KVExpire count/ttl = %d/%v, want 1/%v", client.expireCount, client.lastExpireTTL, sessionIDTTL)
	}
	if client.lastSetTTL != sessionIDTTL {
		t.Fatalf("KVSetNX ttl = %v, want %v", client.lastSetTTL, sessionIDTTL)
	}
}

func TestCachedSessionIDRequiredEmptyAPIKeyDoesNotUseHomeKV(t *testing.T) {
	client := newFakeClaudeIDKVClient()
	useFakeClaudeIDKVClient(t, client, true, nil)

	value, errValue := CachedSessionIDRequired(context.Background(), "")
	if errValue != nil {
		t.Fatalf("CachedSessionIDRequired(empty) error = %v", errValue)
	}
	if _, errParse := uuid.Parse(value); errParse != nil {
		t.Fatalf("session id %q is not a UUID: %v", value, errParse)
	}
	if client.getCount != 0 || client.setCount != 0 || client.expireCount != 0 {
		t.Fatalf("KV calls = get %d set %d expire %d, want all zero", client.getCount, client.setCount, client.expireCount)
	}
}

func TestCachedSessionIDRequiredHomeKVFailures(t *testing.T) {
	for _, tc := range []struct {
		name   string
		client *fakeClaudeIDKVClient
	}{
		{name: "get", client: &fakeClaudeIDKVClient{values: make(map[string][]byte), getErr: errors.New("get failed")}},
		{name: "set", client: &fakeClaudeIDKVClient{values: make(map[string][]byte), setErr: errors.New("set failed")}},
		{name: "expire", client: &fakeClaudeIDKVClient{values: map[string][]byte{
			claudeSessionIDKVKey("api-key-1"): []byte(uuid.New().String()),
		}, expireErr: errors.New("expire failed")}},
	} {
		t.Run(tc.name, func(t *testing.T) {
			useFakeClaudeIDKVClient(t, tc.client, true, nil)
			if _, errValue := CachedSessionIDRequired(context.Background(), "api-key-1"); errValue == nil {
				t.Fatalf("CachedSessionIDRequired() error = nil, want error")
			}
		})
	}
}

func TestCachedSessionIDRequiredHomeRequiresReadAfterSet(t *testing.T) {
	client := newFakeClaudeIDKVClient()
	client.setNoPersist = true
	useFakeClaudeIDKVClient(t, client, true, nil)

	if _, errValue := CachedSessionIDRequired(context.Background(), "api-key-1"); errValue == nil {
		t.Fatalf("CachedSessionIDRequired() error = nil, want missing-after-set error")
	}
}

func TestCachedSessionIDRequiredNonHomeModeUsesLocalMap(t *testing.T) {
	resetSessionIDCache()
	client := newFakeClaudeIDKVClient()
	useFakeClaudeIDKVClient(t, client, false, nil)

	first, errFirst := CachedSessionIDRequired(context.Background(), "api-key-1")
	if errFirst != nil {
		t.Fatalf("CachedSessionIDRequired() first error = %v", errFirst)
	}
	second, errSecond := CachedSessionIDRequired(context.Background(), "api-key-1")
	if errSecond != nil {
		t.Fatalf("CachedSessionIDRequired() second error = %v", errSecond)
	}
	if first != second {
		t.Fatalf("session id = %q then %q, want local cache reuse", first, second)
	}
	if client.getCount != 0 || client.setCount != 0 || client.expireCount != 0 {
		t.Fatalf("KV calls = get %d set %d expire %d, want all zero", client.getCount, client.setCount, client.expireCount)
	}
}