Add scale subresource logic to fake client

kubernetes-sigs · Jun 17, 2024 · 9d4522d · 9d4522d
1 parent 1f5b39f
commit 9d4522d
Show file tree

Hide file tree

Showing 2 changed files with 174 additions and 4 deletions.
diff --git a/pkg/client/fake/client.go b/pkg/client/fake/client.go
@@ -31,6 +31,8 @@ import (
 
 	// Using v4 to match upstream
 	jsonpatch "github.com/evanphx/json-patch"
+	appsv1 "k8s.io/api/apps/v1"
+	autoscalingv1 "k8s.io/api/autoscaling/v1"
 	corev1 "k8s.io/api/core/v1"
 	policyv1 "k8s.io/api/policy/v1"
 	policyv1beta1 "k8s.io/api/policy/v1beta1"
@@ -1080,7 +1082,24 @@ type fakeSubResourceClient struct {
 }
 
 func (sw *fakeSubResourceClient) Get(ctx context.Context, obj, subResource client.Object, opts ...client.SubResourceGetOption) error {
-	panic("fakeSubResourceClient does not support get")
+	switch sw.subResource {
+	case "scale":
+		scale, isScale := subResource.(*autoscalingv1.Scale)
+		if !isScale {
+			return apierrors.NewBadRequest(fmt.Sprintf("got invalid type %t, expected Scale", subResource))
+		}
+		if err := sw.client.Get(ctx, client.ObjectKeyFromObject(obj), obj); err != nil {
+			return err
+		}
+		scaleOut, err := extractScale(obj)
+		if err != nil {
+			return err
+		}
+		*scale = scaleOut
+		return nil
+	default:
+		return fmt.Errorf("fakeSubResourceClient does not support get for %s", sw.subResource)
+	}
 }
 
 func (sw *fakeSubResourceClient) Create(ctx context.Context, obj client.Object, subResource client.Object, opts ...client.SubResourceCreateOption) error {
@@ -1108,10 +1127,28 @@ func (sw *fakeSubResourceClient) Update(ctx context.Context, obj client.Object,
 	updateOptions.ApplyOptions(opts)
 
 	body := obj
-	if updateOptions.SubResourceBody != nil {
-		body = updateOptions.SubResourceBody
+	switch sw.subResource {
+	case "scale":
+		if updateOptions.SubResourceBody == nil {
+			return apierrors.NewBadRequest("expected SubResourceBody")
+		}
+		scale, isScale := updateOptions.SubResourceBody.(*autoscalingv1.Scale)
+		if !isScale {
+			return apierrors.NewBadRequest(fmt.Sprintf("got invalid type %t, expected Scale", updateOptions.SubResourceBody))
+		}
+		if err := sw.client.Get(ctx, client.ObjectKeyFromObject(obj), obj); err != nil {
+			return err
+		}
+		if err := applyScale(obj, scale); err != nil {
+			return err
+		}
+		return sw.client.update(body, false, &updateOptions.UpdateOptions)
+	default:
+		if updateOptions.SubResourceBody != nil {
+			body = updateOptions.SubResourceBody
+		}
+		return sw.client.update(body, true, &updateOptions.UpdateOptions)
 	}
-	return sw.client.update(body, true, &updateOptions.UpdateOptions)
 }
 
 func (sw *fakeSubResourceClient) Patch(ctx context.Context, obj client.Object, patch client.Patch, opts ...client.SubResourcePatchOption) error {
@@ -1278,3 +1315,88 @@ func zero(x interface{}) {
 	res := reflect.ValueOf(x).Elem()
 	res.Set(reflect.Zero(res.Type()))
 }
+
+func extractScale(obj client.Object) (autoscalingv1.Scale, error) {
+	switch obj := obj.(type) {
+	case *appsv1.Deployment:
+		var replicas int32 = 1
+		if obj.Spec.Replicas != nil {
+			replicas = *obj.Spec.Replicas
+		}
+		return autoscalingv1.Scale{
+			ObjectMeta: obj.ObjectMeta,
+			Spec: autoscalingv1.ScaleSpec{
+				Replicas: replicas,
+			},
+			Status: autoscalingv1.ScaleStatus{
+				Replicas: obj.Status.Replicas,
+				Selector: obj.Spec.Selector.String(),
+			},
+		}, nil
+	case *appsv1.ReplicaSet:
+		var replicas int32 = 1
+		if obj.Spec.Replicas != nil {
+			replicas = *obj.Spec.Replicas
+		}
+		return autoscalingv1.Scale{
+			ObjectMeta: obj.ObjectMeta,
+			Spec: autoscalingv1.ScaleSpec{
+				Replicas: replicas,
+			},
+			Status: autoscalingv1.ScaleStatus{
+				Replicas: obj.Status.Replicas,
+				Selector: obj.Spec.Selector.String(),
+			},
+		}, nil
+	case *corev1.ReplicationController:
+		var replicas int32 = 1
+		if obj.Spec.Replicas != nil {
+			replicas = *obj.Spec.Replicas
+		}
+		return autoscalingv1.Scale{
+			ObjectMeta: obj.ObjectMeta,
+			Spec: autoscalingv1.ScaleSpec{
+				Replicas: replicas,
+			},
+			Status: autoscalingv1.ScaleStatus{
+				Replicas: obj.Status.Replicas,
+				Selector: labels.Set(obj.Spec.Selector).String(),
+			},
+		}, nil
+	case *appsv1.StatefulSet:
+		var replicas int32 = 1
+		if obj.Spec.Replicas != nil {
+			replicas = *obj.Spec.Replicas
+		}
+		return autoscalingv1.Scale{
+			ObjectMeta: obj.ObjectMeta,
+			Spec: autoscalingv1.ScaleSpec{
+				Replicas: replicas,
+			},
+			Status: autoscalingv1.ScaleStatus{
+				Replicas: obj.Status.Replicas,
+				Selector: obj.Spec.Selector.String(),
+			},
+		}, nil
+	default:
+		// TODO: CRDs https://kubernetes.io/docs/tasks/extend-kubernetes/custom-resources/custom-resource-definitions/#scale-subresource
+		return autoscalingv1.Scale{}, fmt.Errorf("unable to extract scale from type %T", obj)
+	}
+}
+
+func applyScale(obj client.Object, scale *autoscalingv1.Scale) error {
+	switch obj := obj.(type) {
+	case *appsv1.Deployment:
+		obj.Spec.Replicas = &scale.Spec.Replicas
+	case *appsv1.ReplicaSet:
+		obj.Spec.Replicas = &scale.Spec.Replicas
+	case *corev1.ReplicationController:
+		obj.Spec.Replicas = &scale.Spec.Replicas
+	case *appsv1.StatefulSet:
+		obj.Spec.Replicas = &scale.Spec.Replicas
+	default:
+		// TODO: CRDs https://kubernetes.io/docs/tasks/extend-kubernetes/custom-resources/custom-resource-definitions/#scale-subresource
+		return fmt.Errorf("unable to extract scale from type %T", obj)
+	}
+	return nil
+}
diff --git a/pkg/client/fake/client_test.go b/pkg/client/fake/client_test.go
@@ -27,6 +27,7 @@ import (
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
 	appsv1 "k8s.io/api/apps/v1"
+	autoscalingv1 "k8s.io/api/autoscaling/v1"
 	coordinationv1 "k8s.io/api/coordination/v1"
 	corev1 "k8s.io/api/core/v1"
 	policyv1 "k8s.io/api/policy/v1"
@@ -2068,6 +2069,53 @@ var _ = Describe("Fake client", func() {
 		err := cl.Get(context.Background(), client.ObjectKey{Name: "foo"}, obj)
 		Expect(apierrors.IsNotFound(err)).To(BeTrue())
 	})
+
+	It("should be able to Get scale subresources", func() {
+		obj := &appsv1.Deployment{
+			ObjectMeta: metav1.ObjectMeta{
+				Name: "deploy",
+			},
+			Spec: appsv1.DeploymentSpec{
+				Replicas: ptr.To[int32](2),
+			},
+		}
+		cl := NewClientBuilder().WithObjects(obj).Build()
+		objOriginal := obj.DeepCopy()
+
+		scale := &autoscalingv1.Scale{}
+		Expect(cl.SubResource("scale").Get(context.Background(), obj, scale)).NotTo(HaveOccurred())
+
+		actual := &appsv1.Deployment{ObjectMeta: metav1.ObjectMeta{Name: obj.Name}}
+		Expect(cl.Get(context.Background(), client.ObjectKeyFromObject(actual), actual)).To(Succeed())
+
+		objOriginal.APIVersion = scale.APIVersion
+		objOriginal.Kind = scale.Kind
+		objOriginal.ResourceVersion = scale.ResourceVersion
+		objOriginal.Spec.Replicas = ptr.To(scale.Spec.Replicas)
+		Expect(cmp.Diff(objOriginal, actual)).To(BeEmpty())
+	})
+
+	It("should be able to Update scale subresources", func() {
+		obj := &appsv1.Deployment{
+			ObjectMeta: metav1.ObjectMeta{
+				Name: "deploy",
+			},
+		}
+		cl := NewClientBuilder().WithObjects(obj).Build()
+		objOriginal := obj.DeepCopy()
+
+		scale := &autoscalingv1.Scale{Spec: autoscalingv1.ScaleSpec{Replicas: 2}}
+		Expect(cl.SubResource("scale").Update(context.Background(), obj, client.WithSubResourceBody(scale))).NotTo(HaveOccurred())
+
+		actual := &appsv1.Deployment{ObjectMeta: metav1.ObjectMeta{Name: obj.Name}}
+		Expect(cl.Get(context.Background(), client.ObjectKeyFromObject(actual), actual)).To(Succeed())
+
+		objOriginal.APIVersion = actual.APIVersion
+		objOriginal.Kind = actual.Kind
+		objOriginal.ResourceVersion = actual.ResourceVersion
+		objOriginal.Spec.Replicas = ptr.To(int32(2))
+		Expect(cmp.Diff(objOriginal, actual)).To(BeEmpty())
+	})
 })
 
 type WithPointerMetaList struct {