From c99280d3c7b9bf3132276b76e6acf4dfb500f053 Mon Sep 17 00:00:00 2001
From: Jack Francis <jackfrancis@gmail.com>
Date: Tue, 15 Apr 2025 16:33:40 -0700
Subject: [PATCH] Allow draining when StatefulSet kind has custom API Group

Signed-off-by: Jack Francis <jackfrancis@gmail.com>
---
 .../drainability/rules/replicacount/rule.go   |  4 ++++
 .../rules/replicated/rule_test.go             | 24 +++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/cluster-autoscaler/simulator/drainability/rules/replicacount/rule.go b/cluster-autoscaler/simulator/drainability/rules/replicacount/rule.go
index e03c6d365e..15d408dd7c 100644
--- a/cluster-autoscaler/simulator/drainability/rules/replicacount/rule.go
+++ b/cluster-autoscaler/simulator/drainability/rules/replicacount/rule.go
@@ -110,6 +110,10 @@ func (r *Rule) Drainable(drainCtx *drainability.DrainContext, pod *apiv1.Pod) dr
 			return drainability.NewBlockedStatus(drain.ControllerNotFound, fmt.Errorf("replication controller for %s/%s is not available, err: %v", pod.Namespace, pod.Name, err))
 		}
 	} else if refKind == "StatefulSet" {
+		if refGroup.Group != "apps" {
+			// We don't have a listener for the other StatefulSet group.
+			return drainability.NewUndefinedStatus()
+		}
 		ss, err := drainCtx.Listers.StatefulSetLister().StatefulSets(controllerNamespace).Get(controllerRef.Name)
 
 		if err != nil && ss == nil {
diff --git a/cluster-autoscaler/simulator/drainability/rules/replicated/rule_test.go b/cluster-autoscaler/simulator/drainability/rules/replicated/rule_test.go
index ce2943eb42..44c3dd3a09 100644
--- a/cluster-autoscaler/simulator/drainability/rules/replicated/rule_test.go
+++ b/cluster-autoscaler/simulator/drainability/rules/replicated/rule_test.go
@@ -143,6 +143,30 @@ func TestDrainable(t *testing.T) {
 			},
 			rcs: []*apiv1.ReplicationController{&rc},
 		},
+		"SS-managed pod by a custom API Group": {
+			pod: &apiv1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:            "bar",
+					Namespace:       "default",
+					OwnerReferences: test.GenerateOwnerReferences(statefulset.Name, "StatefulSet", "kruise/v1", ""),
+				},
+				Spec: apiv1.PodSpec{
+					NodeName: "node",
+				},
+			},
+		},
+		"SS-managed pod by a custom API Group with missing reference": {
+			pod: &apiv1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:            "bar",
+					Namespace:       "default",
+					OwnerReferences: test.GenerateOwnerReferences("missing", "StatefulSet", "kruise/v1", ""),
+				},
+				Spec: apiv1.PodSpec{
+					NodeName: "node",
+				},
+			},
+		},
 		"RS-managed pod": {
 			pod: &apiv1.Pod{
 				ObjectMeta: metav1.ObjectMeta{