fix: Redis index parsing bugs and add kvblock unit tests (#80)

Signed-off-by: Kay Yan <[email protected]>

Author: yankay

go.mod

@@ -3,6 +3,7 @@ module github.com/llm-d/llm-d-kv-cache-manager
 go 1.24.1
 
 require (
+	github.com/alicebob/miniredis/v2 v2.35.0
 	github.com/cespare/xxhash/v2 v2.3.0
 	github.com/daulet/tokenizers v1.22.1
 	github.com/fxamacker/cbor/v2 v2.7.0
@@ -44,6 +45,7 @@ require (
 	github.com/prometheus/procfs v0.15.1 // indirect
 	github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect
 	github.com/x448/float16 v0.8.4 // indirect
+	github.com/yuin/gopher-lua v1.1.1 // indirect
 	golang.org/x/net v0.38.0 // indirect
 	golang.org/x/oauth2 v0.27.0 // indirect
 	golang.org/x/sys v0.31.0 // indirect

go.sum

@@ -1,3 +1,5 @@
+github.com/alicebob/miniredis/v2 v2.35.0 h1:QwLphYqCEAo1eu1TqPRN2jgVMPBweeQcR21jeqDCONI=
+github.com/alicebob/miniredis/v2 v2.35.0/go.mod h1:TcL7YfarKPGDAthEtl5NBeHZfeUQj6OXMm/+iu5cLMM=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
 github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
 github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs=
@@ -111,6 +113,8 @@ github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM=
 github.com/x448/float16 v0.8.4/go.mod h1:14CWIYCyZA/cWjXOioeEpHeN/83MdbZDRQHoFcYsOfg=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
+github.com/yuin/gopher-lua v1.1.1 h1:kYKnWBjvbNP4XLT3+bPEwAXJx262OhaHDWDVOPjL46M=
+github.com/yuin/gopher-lua v1.1.1/go.mod h1:GBR0iDaNXjAgGg9zfCvksxSRnQx76gclCIb7kdAd1Pw=
 go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
 go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=

pkg/kvcache/kvblock/common_test.go

@@ -0,0 +1,47 @@
+/*
+Copyright 2025 The llm-d Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package kvblock_test
+
+import (
+	"testing"
+
+	"github.com/llm-d/llm-d-kv-cache-manager/pkg/kvcache/kvblock"
+	"github.com/stretchr/testify/assert"
+	"k8s.io/apimachinery/pkg/util/sets"
+)
+
+// testAddBasic is a common test helper function for testing basic Add and Lookup functionality.
+func testAddBasic(t *testing.T, index kvblock.Index) {
+	t.Helper()
+
+	key := kvblock.Key{ModelName: "12345", ChunkHash: 12345}
+	entries := []kvblock.PodEntry{
+		{PodIdentifier: "10.0.0.1", DeviceTier: "gpu"},
+		{PodIdentifier: "10.0.0.2", DeviceTier: "gpu"},
+	}
+
+	// Add entries
+	err := index.Add(t.Context(), []kvblock.Key{key}, entries)
+	assert.NoError(t, err)
+
+	// Lookup after add
+	hitKeys, podsPerKey, err := index.Lookup(t.Context(), []kvblock.Key{key}, sets.Set[string]{})
+	assert.NoError(t, err)
+	assert.Len(t, hitKeys, 1)
+	assert.Equal(t, key, hitKeys[0])
+	assert.Equal(t, podsPerKey[key], []string{"10.0.0.1", "10.0.0.2"})
+}

pkg/kvcache/kvblock/in_memory_test.go

@@ -0,0 +1,31 @@
+/*
+Copyright 2025 The llm-d Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package kvblock_test
+
+import (
+	"testing"
+
+	"github.com/llm-d/llm-d-kv-cache-manager/pkg/kvcache/kvblock"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestInMemoryAddBasic(t *testing.T) {
+	// Create index
+	index, err := kvblock.NewInMemoryIndex(nil)
+	assert.NoError(t, err)
+	testAddBasic(t, index)
+}

pkg/kvcache/kvblock/redis.go

@@ -126,7 +126,7 @@ func (r *RedisIndex) Lookup(ctx context.Context, keys []Key,
 
 		var filteredPods []string
 		for _, p := range pods {
-			ip := strings.SplitN(p, ":", 2)[0]
+			ip := strings.SplitN(p, "@", 2)[0]
 			if !filterPods || podIdentifierSet.Has(ip) {
 				filteredPods = append(filteredPods, ip)
 			}
@@ -141,7 +141,7 @@ func (r *RedisIndex) Lookup(ctx context.Context, keys []Key,
 		podsPerKey[key] = filteredPods
 	}
 
-	return keys[:highestHitIdx], podsPerKey, nil
+	return keys[:highestHitIdx+1], podsPerKey, nil
 }
 
 // Add adds a set of keys and their associated pod entries to the index backend.

pkg/kvcache/kvblock/redis_test.go

@@ -0,0 +1,42 @@
+/*
+Copyright 2025 The llm-d Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package kvblock_test
+
+import (
+	"testing"
+
+	"github.com/alicebob/miniredis/v2"
+	"github.com/llm-d/llm-d-kv-cache-manager/pkg/kvcache/kvblock"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRedisAddBasic(t *testing.T) {
+	// Create index
+	server, err := miniredis.Run()
+	if err != nil {
+		t.Fatalf("failed to start miniredis: %v", err)
+	}
+	defer server.Close()
+
+	redisConfig := &kvblock.RedisIndexConfig{
+		Address: server.Addr(),
+	}
+	index, err := kvblock.NewRedisIndex(redisConfig)
+	assert.NoError(t, err)
+
+	testAddBasic(t, index)
+}