513 lines
14 KiB
Go
513 lines
14 KiB
Go
/*
|
|
* SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
*
|
|
* 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 controller
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/ai-dynamo/dynamo/deploy/operator/api/v1alpha1"
|
|
"github.com/ai-dynamo/dynamo/deploy/operator/internal/consts"
|
|
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
|
|
"k8s.io/apimachinery/pkg/types"
|
|
"k8s.io/client-go/kubernetes/scheme"
|
|
"k8s.io/client-go/tools/record"
|
|
"k8s.io/utils/ptr"
|
|
ctrl "sigs.k8s.io/controller-runtime"
|
|
"sigs.k8s.io/controller-runtime/pkg/client"
|
|
"sigs.k8s.io/controller-runtime/pkg/client/fake"
|
|
)
|
|
|
|
func TestDynamoGraphDeploymentScalingAdapterReconciler_Reconcile(t *testing.T) {
|
|
// Register custom types with the scheme
|
|
if err := v1alpha1.AddToScheme(scheme.Scheme); err != nil {
|
|
t.Fatalf("Failed to add v1alpha1 to scheme: %v", err)
|
|
}
|
|
|
|
tests := []struct {
|
|
name string
|
|
adapter *v1alpha1.DynamoGraphDeploymentScalingAdapter
|
|
dgd *v1alpha1.DynamoGraphDeployment
|
|
expectedDGDReplicas int32
|
|
expectedStatusReplicas int32
|
|
expectError bool
|
|
expectRequeue bool
|
|
}{
|
|
{
|
|
name: "updates DGD replicas when DGDSA spec differs",
|
|
adapter: &v1alpha1.DynamoGraphDeploymentScalingAdapter{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test-dgd-frontend",
|
|
Namespace: "default",
|
|
},
|
|
Spec: v1alpha1.DynamoGraphDeploymentScalingAdapterSpec{
|
|
Replicas: 5,
|
|
DGDRef: v1alpha1.DynamoGraphDeploymentServiceRef{
|
|
Name: "test-dgd",
|
|
ServiceName: "Frontend",
|
|
},
|
|
},
|
|
},
|
|
dgd: &v1alpha1.DynamoGraphDeployment{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test-dgd",
|
|
Namespace: "default",
|
|
},
|
|
Spec: v1alpha1.DynamoGraphDeploymentSpec{
|
|
Services: map[string]*v1alpha1.DynamoComponentDeploymentSharedSpec{
|
|
"Frontend": {
|
|
Replicas: ptr.To(int32(2)),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
expectedDGDReplicas: 5,
|
|
expectedStatusReplicas: 5,
|
|
expectError: false,
|
|
},
|
|
{
|
|
name: "no update when replicas already match",
|
|
adapter: &v1alpha1.DynamoGraphDeploymentScalingAdapter{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test-dgd-frontend",
|
|
Namespace: "default",
|
|
},
|
|
Spec: v1alpha1.DynamoGraphDeploymentScalingAdapterSpec{
|
|
Replicas: 3,
|
|
DGDRef: v1alpha1.DynamoGraphDeploymentServiceRef{
|
|
Name: "test-dgd",
|
|
ServiceName: "Frontend",
|
|
},
|
|
},
|
|
},
|
|
dgd: &v1alpha1.DynamoGraphDeployment{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test-dgd",
|
|
Namespace: "default",
|
|
},
|
|
Spec: v1alpha1.DynamoGraphDeploymentSpec{
|
|
Services: map[string]*v1alpha1.DynamoComponentDeploymentSharedSpec{
|
|
"Frontend": {
|
|
Replicas: ptr.To(int32(3)),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
expectedDGDReplicas: 3,
|
|
expectedStatusReplicas: 3,
|
|
expectError: false,
|
|
},
|
|
{
|
|
name: "uses default replicas (1) when DGD service has no replicas set",
|
|
adapter: &v1alpha1.DynamoGraphDeploymentScalingAdapter{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test-dgd-worker",
|
|
Namespace: "default",
|
|
},
|
|
Spec: v1alpha1.DynamoGraphDeploymentScalingAdapterSpec{
|
|
Replicas: 4,
|
|
DGDRef: v1alpha1.DynamoGraphDeploymentServiceRef{
|
|
Name: "test-dgd",
|
|
ServiceName: "worker",
|
|
},
|
|
},
|
|
},
|
|
dgd: &v1alpha1.DynamoGraphDeployment{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test-dgd",
|
|
Namespace: "default",
|
|
},
|
|
Spec: v1alpha1.DynamoGraphDeploymentSpec{
|
|
Services: map[string]*v1alpha1.DynamoComponentDeploymentSharedSpec{
|
|
"worker": {}, // no replicas set
|
|
},
|
|
},
|
|
},
|
|
expectedDGDReplicas: 4,
|
|
expectedStatusReplicas: 4,
|
|
expectError: false,
|
|
},
|
|
{
|
|
name: "error when service not found in DGD",
|
|
adapter: &v1alpha1.DynamoGraphDeploymentScalingAdapter{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test-dgd-missing",
|
|
Namespace: "default",
|
|
},
|
|
Spec: v1alpha1.DynamoGraphDeploymentScalingAdapterSpec{
|
|
Replicas: 2,
|
|
DGDRef: v1alpha1.DynamoGraphDeploymentServiceRef{
|
|
Name: "test-dgd",
|
|
ServiceName: "nonexistent",
|
|
},
|
|
},
|
|
},
|
|
dgd: &v1alpha1.DynamoGraphDeployment{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test-dgd",
|
|
Namespace: "default",
|
|
},
|
|
Spec: v1alpha1.DynamoGraphDeploymentSpec{
|
|
Services: map[string]*v1alpha1.DynamoComponentDeploymentSharedSpec{
|
|
"Frontend": {
|
|
Replicas: ptr.To(int32(1)),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
expectError: true,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
// Build initial objects
|
|
var initObjs []client.Object
|
|
initObjs = append(initObjs, tt.adapter, tt.dgd)
|
|
|
|
// Create fake client with status subresource support
|
|
fakeClient := fake.NewClientBuilder().
|
|
WithScheme(scheme.Scheme).
|
|
WithObjects(initObjs...).
|
|
WithStatusSubresource(&v1alpha1.DynamoGraphDeploymentScalingAdapter{}).
|
|
Build()
|
|
|
|
// Create reconciler
|
|
r := &DynamoGraphDeploymentScalingAdapterReconciler{
|
|
Client: fakeClient,
|
|
Scheme: scheme.Scheme,
|
|
Recorder: record.NewFakeRecorder(10),
|
|
}
|
|
|
|
// Run Reconcile
|
|
ctx := context.Background()
|
|
req := ctrl.Request{
|
|
NamespacedName: types.NamespacedName{
|
|
Name: tt.adapter.Name,
|
|
Namespace: tt.adapter.Namespace,
|
|
},
|
|
}
|
|
|
|
result, err := r.Reconcile(ctx, req)
|
|
|
|
// Check error expectation
|
|
if tt.expectError && err == nil {
|
|
t.Errorf("Expected error, but got none")
|
|
}
|
|
if !tt.expectError && err != nil {
|
|
t.Errorf("Unexpected error: %v", err)
|
|
}
|
|
|
|
// Skip further checks if error was expected
|
|
if tt.expectError {
|
|
return
|
|
}
|
|
|
|
// Check requeue
|
|
if tt.expectRequeue && result.RequeueAfter == 0 {
|
|
t.Errorf("Expected requeue, but got none")
|
|
}
|
|
|
|
// Verify DGD replicas were updated
|
|
updatedDGD := &v1alpha1.DynamoGraphDeployment{}
|
|
if err := fakeClient.Get(ctx, types.NamespacedName{Name: tt.dgd.Name, Namespace: tt.dgd.Namespace}, updatedDGD); err != nil {
|
|
t.Fatalf("Failed to get updated DGD: %v", err)
|
|
}
|
|
|
|
service, exists := updatedDGD.Spec.Services[tt.adapter.Spec.DGDRef.ServiceName]
|
|
if !exists {
|
|
t.Fatalf("Service %s not found in updated DGD", tt.adapter.Spec.DGDRef.ServiceName)
|
|
}
|
|
|
|
actualReplicas := int32(1)
|
|
if service.Replicas != nil {
|
|
actualReplicas = *service.Replicas
|
|
}
|
|
|
|
if actualReplicas != tt.expectedDGDReplicas {
|
|
t.Errorf("DGD service replicas = %d, expected %d", actualReplicas, tt.expectedDGDReplicas)
|
|
}
|
|
|
|
// Verify adapter status was updated
|
|
updatedAdapter := &v1alpha1.DynamoGraphDeploymentScalingAdapter{}
|
|
if err := fakeClient.Get(ctx, types.NamespacedName{Name: tt.adapter.Name, Namespace: tt.adapter.Namespace}, updatedAdapter); err != nil {
|
|
t.Fatalf("Failed to get updated adapter: %v", err)
|
|
}
|
|
|
|
if updatedAdapter.Status.Replicas != tt.expectedStatusReplicas {
|
|
t.Errorf("Adapter status.replicas = %d, expected %d", updatedAdapter.Status.Replicas, tt.expectedStatusReplicas)
|
|
}
|
|
|
|
// Verify selector is set
|
|
if updatedAdapter.Status.Selector == "" {
|
|
t.Errorf("Adapter status.selector is empty, expected non-empty")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestDynamoGraphDeploymentScalingAdapterReconciler_Reconcile_NotFound(t *testing.T) {
|
|
// Register custom types with the scheme
|
|
if err := v1alpha1.AddToScheme(scheme.Scheme); err != nil {
|
|
t.Fatalf("Failed to add v1alpha1 to scheme: %v", err)
|
|
}
|
|
|
|
// Create fake client with no objects
|
|
fakeClient := fake.NewClientBuilder().
|
|
WithScheme(scheme.Scheme).
|
|
Build()
|
|
|
|
r := &DynamoGraphDeploymentScalingAdapterReconciler{
|
|
Client: fakeClient,
|
|
Scheme: scheme.Scheme,
|
|
Recorder: record.NewFakeRecorder(10),
|
|
}
|
|
|
|
ctx := context.Background()
|
|
req := ctrl.Request{
|
|
NamespacedName: types.NamespacedName{
|
|
Name: "nonexistent",
|
|
Namespace: "default",
|
|
},
|
|
}
|
|
|
|
// Should return no error when adapter not found (client.IgnoreNotFound)
|
|
result, err := r.Reconcile(ctx, req)
|
|
if err != nil {
|
|
t.Errorf("Expected no error for not found adapter, got: %v", err)
|
|
}
|
|
if result.RequeueAfter != 0 {
|
|
t.Errorf("Expected no requeueAfter for not found adapter, got: %v", result.RequeueAfter)
|
|
}
|
|
}
|
|
|
|
func TestDynamoGraphDeploymentScalingAdapterReconciler_Reconcile_DGDNotFound(t *testing.T) {
|
|
// Register custom types with the scheme
|
|
if err := v1alpha1.AddToScheme(scheme.Scheme); err != nil {
|
|
t.Fatalf("Failed to add v1alpha1 to scheme: %v", err)
|
|
}
|
|
|
|
adapter := &v1alpha1.DynamoGraphDeploymentScalingAdapter{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test-dgd-frontend",
|
|
Namespace: "default",
|
|
},
|
|
Spec: v1alpha1.DynamoGraphDeploymentScalingAdapterSpec{
|
|
Replicas: 5,
|
|
DGDRef: v1alpha1.DynamoGraphDeploymentServiceRef{
|
|
Name: "nonexistent-dgd",
|
|
ServiceName: "Frontend",
|
|
},
|
|
},
|
|
}
|
|
|
|
fakeClient := fake.NewClientBuilder().
|
|
WithScheme(scheme.Scheme).
|
|
WithObjects(adapter).
|
|
Build()
|
|
|
|
r := &DynamoGraphDeploymentScalingAdapterReconciler{
|
|
Client: fakeClient,
|
|
Scheme: scheme.Scheme,
|
|
Recorder: record.NewFakeRecorder(10),
|
|
}
|
|
|
|
ctx := context.Background()
|
|
req := ctrl.Request{
|
|
NamespacedName: types.NamespacedName{
|
|
Name: adapter.Name,
|
|
Namespace: adapter.Namespace,
|
|
},
|
|
}
|
|
|
|
// Should return error when DGD not found
|
|
_, err := r.Reconcile(ctx, req)
|
|
if err == nil {
|
|
t.Errorf("Expected error when DGD not found, got none")
|
|
}
|
|
}
|
|
|
|
func TestDynamoGraphDeploymentScalingAdapterReconciler_Reconcile_BeingDeleted(t *testing.T) {
|
|
// Register custom types with the scheme
|
|
if err := v1alpha1.AddToScheme(scheme.Scheme); err != nil {
|
|
t.Fatalf("Failed to add v1alpha1 to scheme: %v", err)
|
|
}
|
|
|
|
now := metav1.Now()
|
|
adapter := &v1alpha1.DynamoGraphDeploymentScalingAdapter{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test-dgd-frontend",
|
|
Namespace: "default",
|
|
DeletionTimestamp: &now,
|
|
Finalizers: []string{"test-finalizer"}, // Required for deletion timestamp to be set
|
|
},
|
|
Spec: v1alpha1.DynamoGraphDeploymentScalingAdapterSpec{
|
|
Replicas: 5,
|
|
DGDRef: v1alpha1.DynamoGraphDeploymentServiceRef{
|
|
Name: "test-dgd",
|
|
ServiceName: "Frontend",
|
|
},
|
|
},
|
|
}
|
|
|
|
dgd := &v1alpha1.DynamoGraphDeployment{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test-dgd",
|
|
Namespace: "default",
|
|
},
|
|
Spec: v1alpha1.DynamoGraphDeploymentSpec{
|
|
Services: map[string]*v1alpha1.DynamoComponentDeploymentSharedSpec{
|
|
"Frontend": {
|
|
Replicas: ptr.To(int32(2)),
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
fakeClient := fake.NewClientBuilder().
|
|
WithScheme(scheme.Scheme).
|
|
WithObjects(adapter, dgd).
|
|
Build()
|
|
|
|
r := &DynamoGraphDeploymentScalingAdapterReconciler{
|
|
Client: fakeClient,
|
|
Scheme: scheme.Scheme,
|
|
Recorder: record.NewFakeRecorder(10),
|
|
}
|
|
|
|
ctx := context.Background()
|
|
req := ctrl.Request{
|
|
NamespacedName: types.NamespacedName{
|
|
Name: adapter.Name,
|
|
Namespace: adapter.Namespace,
|
|
},
|
|
}
|
|
|
|
// Should return no error and skip reconciliation
|
|
result, err := r.Reconcile(ctx, req)
|
|
if err != nil {
|
|
t.Errorf("Expected no error for deleting adapter, got: %v", err)
|
|
}
|
|
if result.RequeueAfter != 0 {
|
|
t.Errorf("Expected no requeueAfter for deleting adapter, got: %v", result.RequeueAfter)
|
|
}
|
|
|
|
// DGD replicas should NOT be updated (still 2)
|
|
updatedDGD := &v1alpha1.DynamoGraphDeployment{}
|
|
if err := fakeClient.Get(ctx, types.NamespacedName{Name: dgd.Name, Namespace: dgd.Namespace}, updatedDGD); err != nil {
|
|
t.Fatalf("Failed to get DGD: %v", err)
|
|
}
|
|
|
|
if *updatedDGD.Spec.Services["Frontend"].Replicas != 2 {
|
|
t.Errorf("DGD replicas should remain unchanged, got %d", *updatedDGD.Spec.Services["Frontend"].Replicas)
|
|
}
|
|
}
|
|
|
|
func TestDynamoGraphDeploymentScalingAdapterReconciler_findAdaptersForDGD(t *testing.T) {
|
|
// Register custom types with the scheme
|
|
if err := v1alpha1.AddToScheme(scheme.Scheme); err != nil {
|
|
t.Fatalf("Failed to add v1alpha1 to scheme: %v", err)
|
|
}
|
|
|
|
dgd := &v1alpha1.DynamoGraphDeployment{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test-dgd",
|
|
Namespace: "default",
|
|
},
|
|
}
|
|
|
|
// Adapters belonging to test-dgd
|
|
adapter1 := &v1alpha1.DynamoGraphDeploymentScalingAdapter{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test-dgd-frontend",
|
|
Namespace: "default",
|
|
Labels: map[string]string{
|
|
consts.KubeLabelDynamoGraphDeploymentName: "test-dgd",
|
|
},
|
|
},
|
|
Spec: v1alpha1.DynamoGraphDeploymentScalingAdapterSpec{
|
|
DGDRef: v1alpha1.DynamoGraphDeploymentServiceRef{
|
|
Name: "test-dgd",
|
|
ServiceName: "Frontend",
|
|
},
|
|
},
|
|
}
|
|
|
|
adapter2 := &v1alpha1.DynamoGraphDeploymentScalingAdapter{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test-dgd-decode",
|
|
Namespace: "default",
|
|
Labels: map[string]string{
|
|
consts.KubeLabelDynamoGraphDeploymentName: "test-dgd",
|
|
},
|
|
},
|
|
Spec: v1alpha1.DynamoGraphDeploymentScalingAdapterSpec{
|
|
DGDRef: v1alpha1.DynamoGraphDeploymentServiceRef{
|
|
Name: "test-dgd",
|
|
ServiceName: "decode",
|
|
},
|
|
},
|
|
}
|
|
|
|
// Adapter belonging to different DGD
|
|
adapterOther := &v1alpha1.DynamoGraphDeploymentScalingAdapter{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "other-dgd-frontend",
|
|
Namespace: "default",
|
|
Labels: map[string]string{
|
|
consts.KubeLabelDynamoGraphDeploymentName: "other-dgd",
|
|
},
|
|
},
|
|
Spec: v1alpha1.DynamoGraphDeploymentScalingAdapterSpec{
|
|
DGDRef: v1alpha1.DynamoGraphDeploymentServiceRef{
|
|
Name: "other-dgd",
|
|
ServiceName: "Frontend",
|
|
},
|
|
},
|
|
}
|
|
|
|
fakeClient := fake.NewClientBuilder().
|
|
WithScheme(scheme.Scheme).
|
|
WithObjects(adapter1, adapter2, adapterOther).
|
|
Build()
|
|
|
|
r := &DynamoGraphDeploymentScalingAdapterReconciler{
|
|
Client: fakeClient,
|
|
}
|
|
|
|
ctx := context.Background()
|
|
requests := r.findAdaptersForDGD(ctx, dgd)
|
|
|
|
// Should return 2 requests (for test-dgd adapters only)
|
|
if len(requests) != 2 {
|
|
t.Errorf("findAdaptersForDGD() returned %d requests, expected 2", len(requests))
|
|
}
|
|
|
|
// Verify correct adapters are returned
|
|
expectedNames := map[string]bool{
|
|
"test-dgd-frontend": true,
|
|
"test-dgd-decode": true,
|
|
}
|
|
|
|
for _, req := range requests {
|
|
if !expectedNames[req.Name] {
|
|
t.Errorf("Unexpected adapter in results: %s", req.Name)
|
|
}
|
|
}
|
|
}
|