Refactor the code

Signed-off-by: Jiaxin Shan <[email protected]>

Author: Jeffwan

pkg/constants/kvcache.go

@@ -1,3 +1,19 @@
+/*
+Copyright 2025 The Aibrix Team.
+
+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 constants
 
 const (
@@ -13,12 +29,16 @@ const (
 
 	KVCacheAnnotationNodeAffinityDefaultKey = "machine.cluster.vke.volcengine.com/gpu-name"
 
-	// Deprecated: use kvcache backend directly.
-	// distributed, centralized.
+	// This config will be deprecated in future, users should specify kvcache backend directly.
 	KVCacheAnnotationMode              = "kvcache.orchestration.aibrix.ai/mode"
 	KVCacheAnnotationContainerRegistry = "kvcache.orchestration.aibrix.ai/container-registry"
 
 	KVCacheLabelValueRoleCache     = "cache"
 	KVCacheLabelValueRoleMetadata  = "metadata"
 	KVCacheLabelValueRoleKVWatcher = "kvwatcher"
+
+	KVCacheBackendVineyard    = "vineyard"
+	KVCacheBackendHPKV        = "hpkv"
+	KVCacheBackendInfinistore = "infinistore"
+	KVCacheBackendDefault     = KVCacheBackendVineyard
 )

pkg/controller/kvcache/backends/common.go

@@ -18,6 +18,7 @@ package backends
 
 import (
 	"fmt"
+
 	orchestrationv1alpha1 "github.com/vllm-project/aibrix/api/orchestration/v1alpha1"
 	"github.com/vllm-project/aibrix/pkg/constants"
 	corev1 "k8s.io/api/core/v1"

pkg/controller/kvcache/backends/common_test.go

@@ -18,14 +18,15 @@ package backends
 
 import (
 	"fmt"
+	"testing"
+
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	orchestrationv1alpha1 "github.com/vllm-project/aibrix/api/orchestration/v1alpha1"
 	"github.com/vllm-project/aibrix/pkg/constants"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/intstr"
-	"testing"
 )
 
 func TestBuildRedisPod(t *testing.T) {

pkg/controller/kvcache/backends/distributed.go

@@ -20,7 +20,9 @@ import (
 	"context"
 	"errors"
 	"fmt"
+
 	orchestrationv1alpha1 "github.com/vllm-project/aibrix/api/orchestration/v1alpha1"
+	"github.com/vllm-project/aibrix/pkg/constants"
 	"k8s.io/klog/v2"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -39,9 +41,9 @@ func NewDistributedReconciler(c client.Client, backend string) *DistributedRecon
 		BaseReconciler: &BaseReconciler{Client: c},
 	}
 
-	if backend == KVCacheBackendInfinistore {
+	if backend == constants.KVCacheBackendInfinistore {
 		reconciler.Backend = InfiniStoreBackend{}
-	} else if backend == KVCacheBackendHPKV {
+	} else if backend == constants.KVCacheBackendHPKV {
 		reconciler.Backend = HpKVBackend{}
 	} else {
 		panic(fmt.Sprintf("unsupported backend: %s", backend))

pkg/controller/kvcache/backends/distributed_test.go

@@ -0,0 +1,208 @@
+/*
+Copyright 2024 The Aibrix Team.
+
+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 backends
+
+import (
+	"context"
+	"testing"
+
+	appsv1 "k8s.io/api/apps/v1"
+	corev1 "k8s.io/api/core/v1"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/vllm-project/aibrix/api/orchestration/v1alpha1"
+	"github.com/vllm-project/aibrix/pkg/constants"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client/fake"
+)
+
+// -- Test for NewDistributedReconciler --
+
+func TestNewDistributedReconciler(t *testing.T) {
+	c := fake.NewClientBuilder().Build()
+
+	t.Run("infinistore backend", func(t *testing.T) {
+		reconciler := NewDistributedReconciler(c, constants.KVCacheBackendInfinistore)
+		assert.Equal(t, "infinistore", reconciler.Backend.Name())
+	})
+
+	t.Run("hpkv backend", func(t *testing.T) {
+		reconciler := NewDistributedReconciler(c, constants.KVCacheBackendHPKV)
+		assert.Equal(t, "hpkv", reconciler.Backend.Name())
+	})
+
+	t.Run("invalid backend panics", func(t *testing.T) {
+		assert.Panics(t, func() {
+			NewDistributedReconciler(c, "invalid-backend")
+		})
+	})
+}
+
+// -- Test for reconcileMetadataService --
+
+func TestReconcileMetadataService(t *testing.T) {
+	scheme := runtime.NewScheme()
+	_ = corev1.AddToScheme(scheme)
+	c := fake.NewClientBuilder().WithScheme(scheme).Build()
+	r := NewDistributedReconciler(c, constants.KVCacheBackendInfinistore)
+
+	t.Run("fails when neither etcd nor redis configured", func(t *testing.T) {
+		kv := &v1alpha1.KVCache{
+			Spec: v1alpha1.KVCacheSpec{
+				Metadata: &v1alpha1.MetadataConfig{},
+			},
+		}
+		err := r.reconcileMetadataService(context.Background(), kv)
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "either etcd or redis configuration is required")
+	})
+
+	t.Run("succeeds when redis is configured", func(t *testing.T) {
+		kv := &v1alpha1.KVCache{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      "redis-test",
+				Namespace: "default",
+			},
+			Spec: v1alpha1.KVCacheSpec{
+				Metadata: &v1alpha1.MetadataConfig{
+					Redis: &v1alpha1.RedisConfig{
+						Replicas: 1,
+					},
+				},
+			},
+		}
+
+		// Replace backend with mock
+		r.Backend = mockBackend{
+			redis: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "mock-redis",
+					Namespace: "default",
+				},
+			},
+			svc: &corev1.Service{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "mock-redis-service",
+					Namespace: "default",
+				},
+				Spec: corev1.ServiceSpec{
+					Ports: []corev1.ServicePort{
+						{
+							Name:     "redis",
+							Port:     6379,
+							Protocol: corev1.ProtocolTCP,
+						},
+					},
+				},
+			},
+		}
+
+		err := r.reconcileMetadataService(context.Background(), kv)
+		assert.NoError(t, err)
+	})
+}
+
+// -- Test for reconcileRedisService (partial stubbed) --
+
+func TestReconcileRedisService_SingleReplicaAllowed(t *testing.T) {
+	c := fake.NewClientBuilder().Build()
+	r := NewDistributedReconciler(c, constants.KVCacheBackendInfinistore)
+
+	kv := &v1alpha1.KVCache{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "redis-test",
+			Namespace: "default",
+		},
+		Spec: v1alpha1.KVCacheSpec{
+			Metadata: &v1alpha1.MetadataConfig{
+				Redis: &v1alpha1.RedisConfig{
+					Replicas: 1,
+				},
+			},
+		},
+	}
+
+	// Stub backend logic
+	r.Backend = mockBackend{
+		// dummy objects
+		redis: &corev1.Pod{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      "mock-redis",
+				Namespace: "default",
+			},
+		},
+		svc: &corev1.Service{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      "mock-redis-service",
+				Namespace: "default",
+			},
+			Spec: corev1.ServiceSpec{
+				Ports: []corev1.ServicePort{
+					{
+						Name:     "redis",
+						Port:     6379,
+						Protocol: corev1.ProtocolTCP,
+					},
+				},
+			},
+		},
+	}
+
+	err := r.ReconcilePodObject(context.Background(), r.Backend.BuildMetadataPod(kv))
+	assert.NoError(t, err)
+
+	err = r.reconcileRedisService(context.Background(), kv)
+	assert.NoError(t, err)
+}
+
+// -- Mock Backend for isolation --
+
+type mockBackend struct {
+	redis   *corev1.Pod
+	watcher *corev1.Pod
+	svc     *corev1.Service
+	sts     *appsv1.StatefulSet
+}
+
+func (m mockBackend) Name() string {
+	return "mock-backend"
+}
+
+func (m mockBackend) ValidateObject(*v1alpha1.KVCache) error {
+	return nil
+}
+
+func (m mockBackend) BuildMetadataPod(*v1alpha1.KVCache) *corev1.Pod {
+	return m.redis
+}
+
+func (m mockBackend) BuildMetadataService(*v1alpha1.KVCache) *corev1.Service {
+	return m.svc
+}
+
+func (m mockBackend) BuildWatcherPod(*v1alpha1.KVCache) *corev1.Pod {
+	return m.watcher
+}
+
+func (m mockBackend) BuildCacheStatefulSet(*v1alpha1.KVCache) *appsv1.StatefulSet {
+	return m.sts
+}
+
+func (m mockBackend) BuildService(*v1alpha1.KVCache) *corev1.Service {
+	return m.svc
+}