wenshu/k8s-orchestrator/internal/controller/clusteroperation_controller_test.go

package controller

import (
	"context"
	"strings"
	"testing"

	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	ctrlclientfake "sigs.k8s.io/controller-runtime/pkg/client/fake"

	orcv1 "github.com/cyrus/spark-orchestrator/api/v1alpha1"
	"github.com/cyrus/spark-orchestrator/internal/util"
)

func TestEnsureNoRunningPodsOnNodes_Blocking(t *testing.T) {
	scheme := runtime.NewScheme()
	utilruntime.Must(corev1.AddToScheme(scheme))
	utilruntime.Must(orcv1.AddToScheme(scheme))

	cluster := &orcv1.SparkVirtualCluster{
		ObjectMeta: metav1.ObjectMeta{Name: "c1"},
		Status: orcv1.SparkVirtualClusterStatus{
			Namespace: "spark-c1",
		},
	}
	nodeInv := &orcv1.NodeInventory{
		ObjectMeta: metav1.ObjectMeta{Name: "node-a"},
		Status:     orcv1.NodeInventoryStatus{KubernetesNodeName: "k8s-node-a"},
	}
	pod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{Name: "driver-1", Namespace: "spark-c1"},
		Spec:       corev1.PodSpec{NodeName: "k8s-node-a"},
		Status:     corev1.PodStatus{Phase: corev1.PodRunning},
	}

	fakeClient := ctrlclientfake.NewClientBuilder().WithScheme(scheme).WithObjects(cluster, nodeInv, pod).Build()
	r := &ClusterOperationReconciler{Client: fakeClient, Scheme: scheme}

	err := r.ensureNoRunningPodsOnNodes(context.Background(), cluster, []string{"node-a"})
	if err == nil {
		t.Fatalf("expected blocking error, got nil")
	}
}

func TestEnsureNoRunningPodsOnNodes_NoBlockForSucceededPod(t *testing.T) {
	scheme := runtime.NewScheme()
	utilruntime.Must(corev1.AddToScheme(scheme))
	utilruntime.Must(orcv1.AddToScheme(scheme))

	cluster := &orcv1.SparkVirtualCluster{
		ObjectMeta: metav1.ObjectMeta{Name: "c1"},
		Status: orcv1.SparkVirtualClusterStatus{
			Namespace: "spark-c1",
		},
	}
	nodeInv := &orcv1.NodeInventory{
		ObjectMeta: metav1.ObjectMeta{Name: "node-a"},
		Status:     orcv1.NodeInventoryStatus{KubernetesNodeName: "k8s-node-a"},
	}
	pod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{Name: "driver-1", Namespace: "spark-c1"},
		Spec:       corev1.PodSpec{NodeName: "k8s-node-a"},
		Status:     corev1.PodStatus{Phase: corev1.PodSucceeded},
	}

	fakeClient := ctrlclientfake.NewClientBuilder().WithScheme(scheme).WithObjects(cluster, nodeInv, pod).Build()
	r := &ClusterOperationReconciler{Client: fakeClient, Scheme: scheme}

	err := r.ensureNoRunningPodsOnNodes(context.Background(), cluster, []string{"node-a"})
	if err != nil {
		t.Fatalf("expected no error, got %v", err)
	}
}

func TestEnsureNoRunningPodsOnNodes_IgnoreSparkPrePullPod(t *testing.T) {
	scheme := runtime.NewScheme()
	utilruntime.Must(corev1.AddToScheme(scheme))
	utilruntime.Must(orcv1.AddToScheme(scheme))

	cluster := &orcv1.SparkVirtualCluster{
		ObjectMeta: metav1.ObjectMeta{Name: "c1"},
		Status: orcv1.SparkVirtualClusterStatus{
			Namespace: "spark-c1",
		},
	}
	nodeInv := &orcv1.NodeInventory{
		ObjectMeta: metav1.ObjectMeta{Name: "node-a"},
		Status:     orcv1.NodeInventoryStatus{KubernetesNodeName: "k8s-node-a"},
	}
	pod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "spark-prepull-1",
			Namespace: "spark-c1",
			Labels: map[string]string{
				util.LabelManagedBy: util.ManagedByValue,
				util.LabelComponent: util.ComponentSparkPrePull,
			},
		},
		Spec:   corev1.PodSpec{NodeName: "k8s-node-a"},
		Status: corev1.PodStatus{Phase: corev1.PodRunning},
	}

	fakeClient := ctrlclientfake.NewClientBuilder().WithScheme(scheme).WithObjects(cluster, nodeInv, pod).Build()
	r := &ClusterOperationReconciler{Client: fakeClient, Scheme: scheme}

	err := r.ensureNoRunningPodsOnNodes(context.Background(), cluster, []string{"node-a"})
	if err != nil {
		t.Fatalf("expected no error for spark pre-pull pod, got %v", err)
	}
}

func TestEnsureNoRunningPodsOnNodes_IgnoreSparkIvyWarmupPod(t *testing.T) {
	scheme := runtime.NewScheme()
	utilruntime.Must(corev1.AddToScheme(scheme))
	utilruntime.Must(orcv1.AddToScheme(scheme))

	cluster := &orcv1.SparkVirtualCluster{
		ObjectMeta: metav1.ObjectMeta{Name: "c1"},
		Status: orcv1.SparkVirtualClusterStatus{
			Namespace: "spark-c1",
		},
	}
	nodeInv := &orcv1.NodeInventory{
		ObjectMeta: metav1.ObjectMeta{Name: "node-a"},
		Status:     orcv1.NodeInventoryStatus{KubernetesNodeName: "k8s-node-a"},
	}
	pod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "spark-ivy-warmup-1",
			Namespace: "spark-c1",
			Labels: map[string]string{
				util.LabelManagedBy: util.ManagedByValue,
				util.LabelComponent: util.ComponentSparkIvyWarmup,
			},
		},
		Spec:   corev1.PodSpec{NodeName: "k8s-node-a"},
		Status: corev1.PodStatus{Phase: corev1.PodRunning},
	}

	fakeClient := ctrlclientfake.NewClientBuilder().WithScheme(scheme).WithObjects(cluster, nodeInv, pod).Build()
	r := &ClusterOperationReconciler{Client: fakeClient, Scheme: scheme}

	err := r.ensureNoRunningPodsOnNodes(context.Background(), cluster, []string{"node-a"})
	if err != nil {
		t.Fatalf("expected no error for spark ivy warmup pod, got %v", err)
	}
}

func TestHandleScaleInStarRocksClusterSerial(t *testing.T) {
	scheme := runtime.NewScheme()
	utilruntime.Must(corev1.AddToScheme(scheme))
	utilruntime.Must(orcv1.AddToScheme(scheme))

	cluster := &orcv1.StarRocksVirtualCluster{
		ObjectMeta: metav1.ObjectMeta{Name: "sr-1"},
		Spec: orcv1.StarRocksVirtualClusterSpec{
			FeNodeRefs: []string{"fe-1"},
			BeNodeRefs: []string{"be-1", "be-2"},
		},
		Status: orcv1.StarRocksVirtualClusterStatus{
			Phase:           orcv1.StarRocksClusterPhaseReady,
			Namespace:       "starrocks-sr-1",
			AssignedFeNodes: []string{"fe-1"},
			AssignedBeNodes: []string{"be-1", "be-2"},
		},
	}
	node1 := &orcv1.NodeInventory{
		ObjectMeta: metav1.ObjectMeta{Name: "be-1"},
		Status:     orcv1.NodeInventoryStatus{KubernetesNodeName: "k8s-be-1"},
	}
	node2 := &orcv1.NodeInventory{
		ObjectMeta: metav1.ObjectMeta{Name: "be-2"},
		Status:     orcv1.NodeInventoryStatus{KubernetesNodeName: "k8s-be-2"},
	}

	fakeClient := ctrlclientfake.NewClientBuilder().WithScheme(scheme).WithObjects(cluster, node1, node2).Build()
	r := &ClusterOperationReconciler{Client: fakeClient, Scheme: scheme}
	op := &orcv1.ClusterOperation{
		Spec: orcv1.ClusterOperationSpec{
			ResizeStarRocks: &orcv1.ResizeStarRocksClusterSpec{
				ClusterName: "sr-1",
				NodeIDs:     []string{"be-1", "be-2"},
			},
		},
	}

	done, _, _, message, err := r.handleScaleInStarRocksCluster(context.Background(), op)
	if err != nil {
		t.Fatalf("first scale-in step returned error: %v", err)
	}
	if done {
		t.Fatalf("first scale-in step should not complete")
	}
	if !strings.Contains(message, "be-1") {
		t.Fatalf("expected first step to remove be-1, got message: %s", message)
	}

	var updated orcv1.StarRocksVirtualCluster
	if err := r.Get(context.Background(), types.NamespacedName{Name: "sr-1"}, &updated); err != nil {
		t.Fatalf("get updated cluster: %v", err)
	}
	if util.ContainsString(updated.Spec.BeNodeRefs, "be-1") {
		t.Fatalf("be-1 should be removed first")
	}
	if !util.ContainsString(updated.Spec.BeNodeRefs, "be-2") {
		t.Fatalf("be-2 should remain after first serial step")
	}

	updated.Status.AssignedBeNodes = []string{"be-2"}
	if err := r.Update(context.Background(), &updated); err != nil {
		t.Fatalf("update cluster status simulation: %v", err)
	}

	done, _, _, message, err = r.handleScaleInStarRocksCluster(context.Background(), op)
	if err != nil {
		t.Fatalf("second scale-in step returned error: %v", err)
	}
	if done {
		t.Fatalf("second scale-in step should not complete")
	}
	if !strings.Contains(message, "be-2") {
		t.Fatalf("expected second step to remove be-2, got message: %s", message)
	}

	if err := r.Get(context.Background(), types.NamespacedName{Name: "sr-1"}, &updated); err != nil {
		t.Fatalf("get cluster after second step: %v", err)
	}
	if len(updated.Spec.BeNodeRefs) != 0 {
		t.Fatalf("expected all BE nodes removed, got %v", updated.Spec.BeNodeRefs)
	}

	updated.Status.AssignedBeNodes = nil
	if err := r.Update(context.Background(), &updated); err != nil {
		t.Fatalf("update final status simulation: %v", err)
	}

	done, _, _, message, err = r.handleScaleInStarRocksCluster(context.Background(), op)
	if err != nil {
		t.Fatalf("final scale-in step returned error: %v", err)
	}
	if !done {
		t.Fatalf("expected serial scale-in to complete")
	}
	if !strings.Contains(strings.ToLower(message), "completed") {
		t.Fatalf("expected completed message, got: %s", message)
	}
}

func TestHandleDeleteNode_RemovesNodeInventoryAndKubernetesNode(t *testing.T) {
	scheme := runtime.NewScheme()
	utilruntime.Must(corev1.AddToScheme(scheme))
	utilruntime.Must(orcv1.AddToScheme(scheme))

	nodeInv := &orcv1.NodeInventory{
		ObjectMeta: metav1.ObjectMeta{Name: "node-16"},
		Spec:       orcv1.NodeInventorySpec{Address: "172.30.194.144"},
		Status: orcv1.NodeInventoryStatus{
			Phase:              orcv1.NodePhaseIdle,
			KubernetesNodeName: "node-16",
		},
	}
	k8sNode := &corev1.Node{
		ObjectMeta: metav1.ObjectMeta{
			Name: "node-16",
			Labels: map[string]string{
				util.LabelSparkClusterID: "spark-1",
				util.LabelManagedBy:      util.ManagedByValue,
			},
		},
	}
	op := &orcv1.ClusterOperation{
		Spec: orcv1.ClusterOperationSpec{
			DeleteNode: "node-16",
		},
	}

	fakeClient := ctrlclientfake.NewClientBuilder().WithScheme(scheme).WithObjects(nodeInv, k8sNode).Build()
	r := &ClusterOperationReconciler{Client: fakeClient, Scheme: scheme}

	done, _, _, message, err := r.handleDeleteNode(context.Background(), op)
	if err != nil {
		t.Fatalf("handleDeleteNode returned error: %v", err)
	}
	if !done {
		t.Fatalf("expected delete operation completed")
	}
	if !strings.Contains(strings.ToLower(message), "cluster") {
		t.Fatalf("expected cluster removal message, got: %s", message)
	}

	var invAfter orcv1.NodeInventory
	if err := fakeClient.Get(context.Background(), types.NamespacedName{Name: "node-16"}, &invAfter); !apierrors.IsNotFound(err) {
		t.Fatalf("expected node inventory deleted, get err=%v", err)
	}

	var nodeAfter corev1.Node
	if err := fakeClient.Get(context.Background(), types.NamespacedName{Name: "node-16"}, &nodeAfter); !apierrors.IsNotFound(err) {
		t.Fatalf("expected kubernetes node deleted, get err=%v", err)
	}
}

func TestHandleDeleteNode_FallbackDeleteByInternalIP(t *testing.T) {
	scheme := runtime.NewScheme()
	utilruntime.Must(corev1.AddToScheme(scheme))
	utilruntime.Must(orcv1.AddToScheme(scheme))

	nodeInv := &orcv1.NodeInventory{
		ObjectMeta: metav1.ObjectMeta{Name: "node-17"},
		Spec:       orcv1.NodeInventorySpec{Address: "10.0.0.17"},
		Status: orcv1.NodeInventoryStatus{
			Phase: orcv1.NodePhaseIdle,
		},
	}
	k8sNode := &corev1.Node{
		ObjectMeta: metav1.ObjectMeta{Name: "real-node-17"},
		Status: corev1.NodeStatus{
			Addresses: []corev1.NodeAddress{
				{Type: corev1.NodeInternalIP, Address: "10.0.0.17"},
			},
		},
	}
	op := &orcv1.ClusterOperation{
		Spec: orcv1.ClusterOperationSpec{
			DeleteNode: "node-17",
		},
	}

	fakeClient := ctrlclientfake.NewClientBuilder().WithScheme(scheme).WithObjects(nodeInv, k8sNode).Build()
	r := &ClusterOperationReconciler{Client: fakeClient, Scheme: scheme}

	done, _, _, _, err := r.handleDeleteNode(context.Background(), op)
	if err != nil {
		t.Fatalf("handleDeleteNode returned error: %v", err)
	}
	if !done {
		t.Fatalf("expected delete operation completed")
	}

	var nodeAfter corev1.Node
	if err := fakeClient.Get(context.Background(), types.NamespacedName{Name: "real-node-17"}, &nodeAfter); !apierrors.IsNotFound(err) {
		t.Fatalf("expected fallback matched kubernetes node deleted, get err=%v", err)
	}
}

func TestHandleDeleteNode_AllowDeleteProvisioningNodeWithoutClusterRef(t *testing.T) {
	scheme := runtime.NewScheme()
	utilruntime.Must(corev1.AddToScheme(scheme))
	utilruntime.Must(orcv1.AddToScheme(scheme))

	nodeInv := &orcv1.NodeInventory{
		ObjectMeta: metav1.ObjectMeta{Name: "node-24"},
		Spec:       orcv1.NodeInventorySpec{Address: "1.2.3.6"},
		Status: orcv1.NodeInventoryStatus{
			Phase:     orcv1.NodePhaseProvisioning,
			LastError: "dial tcp 1.2.3.6:22: i/o timeout",
		},
	}
	op := &orcv1.ClusterOperation{
		Spec: orcv1.ClusterOperationSpec{
			DeleteNode: "node-24",
		},
	}

	fakeClient := ctrlclientfake.NewClientBuilder().WithScheme(scheme).WithObjects(nodeInv).Build()
	r := &ClusterOperationReconciler{Client: fakeClient, Scheme: scheme}

	done, _, _, _, err := r.handleDeleteNode(context.Background(), op)
	if err != nil {
		t.Fatalf("handleDeleteNode returned error: %v", err)
	}
	if !done {
		t.Fatalf("expected delete operation completed")
	}

	var invAfter orcv1.NodeInventory
	if err := fakeClient.Get(context.Background(), types.NamespacedName{Name: "node-24"}, &invAfter); !apierrors.IsNotFound(err) {
		t.Fatalf("expected node inventory deleted, get err=%v", err)
	}
}

func TestHandleAddNode_FailedPhaseShouldKeepRetrying(t *testing.T) {
	scheme := runtime.NewScheme()
	utilruntime.Must(corev1.AddToScheme(scheme))
	utilruntime.Must(orcv1.AddToScheme(scheme))

	nodeInv := &orcv1.NodeInventory{
		ObjectMeta: metav1.ObjectMeta{Name: "node-30"},
		Status: orcv1.NodeInventoryStatus{
			Phase:     orcv1.NodePhaseFailed,
			LastError: "run join command: Process exited with status 1",
		},
	}
	op := &orcv1.ClusterOperation{
		Spec: orcv1.ClusterOperationSpec{
			AddNode: &orcv1.AddNodeSpec{
				NodeName: "node-30",
				Address:  "10.0.0.30",
			},
		},
	}

	fakeClient := ctrlclientfake.NewClientBuilder().WithScheme(scheme).WithObjects(nodeInv).Build()
	r := &ClusterOperationReconciler{Client: fakeClient, Scheme: scheme}

	done, requeueAfter, resourceRef, message, err := r.handleAddNode(context.Background(), op)
	if err != nil {
		t.Fatalf("handleAddNode returned unexpected error: %v", err)
	}
	if done {
		t.Fatalf("failed phase should keep retrying, not done")
	}
	if resourceRef != "node-30" {
		t.Fatalf("unexpected resourceRef: %s", resourceRef)
	}
	if requeueAfter <= 0 {
		t.Fatalf("expected positive requeueAfter for retry")
	}
	if !strings.Contains(strings.ToLower(message), "retry") {
		t.Fatalf("expected retry message, got: %s", message)
	}
}

func TestTryDeleteOwnedNamespace_DeletesOwnedNamespace(t *testing.T) {
	scheme := runtime.NewScheme()
	utilruntime.Must(corev1.AddToScheme(scheme))
	utilruntime.Must(orcv1.AddToScheme(scheme))

	ns := &corev1.Namespace{
		ObjectMeta: metav1.ObjectMeta{
			Name: "starrocks-sr-9",
			Labels: map[string]string{
				util.LabelManagedBy:           util.ManagedByValue,
				util.LabelNamespaceAutoDelete: "true",
				util.LabelNamespaceOwnerKind:  "starrocks",
				util.LabelNamespaceOwnerName:  "sr-9",
			},
		},
	}
	fakeClient := ctrlclientfake.NewClientBuilder().WithScheme(scheme).WithObjects(ns).Build()
	r := &ClusterOperationReconciler{Client: fakeClient, Scheme: scheme}

	message := r.tryDeleteOwnedNamespace(context.Background(), "starrocks-sr-9", "starrocks", "sr-9")
	if !strings.Contains(strings.ToLower(message), "deletion requested") {
		t.Fatalf("expected namespace deletion request message, got: %s", message)
	}

	var nsAfter corev1.Namespace
	if err := fakeClient.Get(context.Background(), types.NamespacedName{Name: "starrocks-sr-9"}, &nsAfter); !apierrors.IsNotFound(err) {
		t.Fatalf("expected namespace deleted, get err=%v", err)
	}
}

func TestTryDeleteOwnedNamespace_SkipsWhenOwnerMismatch(t *testing.T) {
	scheme := runtime.NewScheme()
	utilruntime.Must(corev1.AddToScheme(scheme))
	utilruntime.Must(orcv1.AddToScheme(scheme))

	ns := &corev1.Namespace{
		ObjectMeta: metav1.ObjectMeta{
			Name: "starrocks-sr-10",
			Labels: map[string]string{
				util.LabelManagedBy:           util.ManagedByValue,
				util.LabelNamespaceAutoDelete: "true",
				util.LabelNamespaceOwnerKind:  "starrocks",
				util.LabelNamespaceOwnerName:  "sr-11",
			},
		},
	}
	fakeClient := ctrlclientfake.NewClientBuilder().WithScheme(scheme).WithObjects(ns).Build()
	r := &ClusterOperationReconciler{Client: fakeClient, Scheme: scheme}

	message := r.tryDeleteOwnedNamespace(context.Background(), "starrocks-sr-10", "starrocks", "sr-10")
	if !strings.Contains(strings.ToLower(message), "owner mismatch") {
		t.Fatalf("expected owner mismatch message, got: %s", message)
	}

	var nsAfter corev1.Namespace
	if err := fakeClient.Get(context.Background(), types.NamespacedName{Name: "starrocks-sr-10"}, &nsAfter); err != nil {
		t.Fatalf("expected namespace remains, get err=%v", err)
	}
}