fix: undo referenced object for workloadRef rollout

- unit test is added

Signed-off-by: Hui Kang <[email protected]>

pkg/kubectl-argo-rollouts/cmd/undo/undo.go

@@ -11,12 +11,15 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/client-go/dynamic"
 	"k8s.io/client-go/kubernetes"
 
 	"github.com/argoproj/argo-rollouts/pkg/apis/rollouts/v1alpha1"
 	"github.com/argoproj/argo-rollouts/pkg/kubectl-argo-rollouts/options"
+	rolloutpkg "github.com/argoproj/argo-rollouts/rollout"
+	routils "github.com/argoproj/argo-rollouts/utils/unstructured"
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	apiequality "k8s.io/apimachinery/pkg/api/equality"
@@ -67,6 +70,8 @@ func NewCmdUndo(o *options.ArgoRolloutsOptions) *cobra.Command {
 // RunUndoRollout performs the execution of 'rollouts undo' sub command
 func RunUndoRollout(rolloutIf dynamic.ResourceInterface, c kubernetes.Interface, name string, toRevision int64) (string, error) {
 	ctx := context.TODO()
+	var err error
+
 	ro, err := rolloutIf.Get(ctx, name, metav1.GetOptions{})
 	if err != nil {
 		return "", err
@@ -93,13 +98,48 @@ func RunUndoRollout(rolloutIf dynamic.ResourceInterface, c kubernetes.Interface,
 		return "", fmt.Errorf("failed restoring revision %d: %v", toRevision, err)
 	}
 
-	// Restore revision
-	if _, err = rolloutIf.Patch(ctx, name, patchType, patch, metav1.PatchOptions{}); err != nil {
+	// Restore revision depending on whether workload ref is nil
+	rollout := routils.ObjectToRollout(ro)
+	if rollout == nil {
+		return "", fmt.Errorf("Invalid rollout object")
+	}
+	if rollout.Spec.WorkloadRef != nil {
+		err = undoWorkloadRef(c, rollout, patchType, patch)
+	} else {
+		_, err = rolloutIf.Patch(ctx, name, patchType, patch, metav1.PatchOptions{})
+	}
+	if err != nil {
 		return "", fmt.Errorf("failed restoring revision %d: %v", toRevision, err)
 	}
 	return fmt.Sprintf("rollout '%s' undo\n", ro.GetName()), nil
 }
 
+func undoWorkloadRef(c kubernetes.Interface, rollout *v1alpha1.Rollout, patchType types.PatchType, patch []byte) error {
+	var err error
+
+	refName := rollout.Spec.WorkloadRef.Name
+	namespace := rollout.GetNamespace()
+
+	gvk := schema.FromAPIVersionAndKind(rollout.Spec.WorkloadRef.APIVersion, rollout.Spec.WorkloadRef.Kind)
+
+	_, ok := rolloutpkg.InfoByGroupKind[gvk.GroupKind()]
+	if !ok {
+		return fmt.Errorf("workload of type %s/%s is not supported", gvk.Group, gvk.Kind)
+	}
+
+	switch rollout.Spec.WorkloadRef.Kind {
+	case "Deployment":
+		_, err = c.AppsV1().Deployments(namespace).Patch(context.TODO(), refName, patchType, patch, metav1.PatchOptions{})
+	case "ReplicaSet":
+		_, err = c.AppsV1().ReplicaSets(namespace).Patch(context.TODO(), refName, patchType, patch, metav1.PatchOptions{})
+	case "PodTemplate":
+		_, err = c.CoreV1().PodTemplates(namespace).Patch(context.TODO(), refName, patchType, patch, metav1.PatchOptions{})
+	default:
+		return fmt.Errorf("Undo workload of type %s/%s is not supported", gvk.Group, gvk.Kind)
+	}
+	return err
+}
+
 func rolloutRevision(ro *unstructured.Unstructured, c kubernetes.Interface, toRevision int64) (*appsv1.ReplicaSet, error) {
 	allRSs, err := getAllReplicaSets(ro, c.AppsV1())
 	if err != nil {

pkg/kubectl-argo-rollouts/cmd/undo/undo_test.go

@@ -2,6 +2,7 @@ package undo
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"testing"
@@ -118,6 +119,34 @@ func TestUndoCmdToRevision(t *testing.T) {
 	assert.Empty(t, stderr)
 }
 
+func TestUndoCmdToRevisionOfWorkloadRef(t *testing.T) {
+	rolloutObjs := testdata.NewCanaryRollout()
+	ro := rolloutObjs.Rollouts[1]
+	ro.Spec.Template = corev1.PodTemplateSpec{
+		ObjectMeta: metav1.ObjectMeta{
+			Name: "test",
+		},
+	}
+	tf, o := options.NewFakeArgoRolloutsOptions(rolloutObjs.AllObjects()...)
+	o.RESTClientGetter = tf.WithNamespace(ro.Namespace)
+	defer tf.Cleanup()
+	cmd := NewCmdUndo(o)
+	cmd.PersistentPreRunE = o.PersistentPreRunE
+	cmd.SetArgs([]string{ro.Name, "--to-revision=29"})
+
+	err := cmd.Execute()
+	assert.Nil(t, err)
+
+	// Verify the current RS has been patched by the oldRS's template
+	currentRs, _ := o.KubeClient.AppsV1().ReplicaSets(ro.Namespace).Get(context.TODO(), "canary-demo-65fb5ffc84", metav1.GetOptions{})
+	assert.Equal(t, "argoproj/rollouts-demo:asdf", currentRs.Spec.Template.Spec.Containers[0].Image)
+
+	stdout := o.Out.(*bytes.Buffer).String()
+	stderr := o.ErrOut.(*bytes.Buffer).String()
+	assert.Equal(t, fmt.Sprintf("rollout '%s' undo\n", ro.Name), stdout)
+	assert.Empty(t, stderr)
+}
+
 func TestUndoCmdToRevisionSkipCurrent(t *testing.T) {
 	rolloutObjs := testdata.NewCanaryRollout()
 	ro := rolloutObjs.Rollouts[0]

pkg/kubectl-argo-rollouts/info/testdata/canary/canary-rollout2.yaml

@@ -0,0 +1,68 @@
+apiVersion: argoproj.io/v1alpha1
+kind: Rollout
+metadata:
+  annotations:
+    rollout.argoproj.io/revision: "31"
+  creationTimestamp: "2019-10-25T06:07:18Z"
+  generation: 429
+  labels:
+    app: canary-demo
+    app.kubernetes.io/instance: jesse-test
+  name: canary-demo-workloadRef
+  namespace: jesse-test
+  resourceVersion: "28253567"
+  selfLink: /apis/argoproj.io/v1alpha1/namespaces/jesse-test/rollouts/canary-demo
+  uid: b350ba76-f6ed-11e9-a15b-42010aa80033
+spec:
+  progressDeadlineSeconds: 30
+  replicas: 5
+  revisionHistoryLimit: 3
+  selector:
+    matchLabels:
+      app: canary-demo
+  strategy:
+    canary:
+      canaryService: canary-demo-preview
+      steps:
+      - setWeight: 20
+      - pause: {}
+      - setWeight: 40
+      - pause:
+          duration: 10s
+      - setWeight: 60
+      - pause:
+          duration: 10s
+      - setWeight: 80
+      - pause:
+          duration: 10s
+  workloadRef:
+    apiVersion: apps/v1
+    kind: ReplicaSet
+    name: canary-demo-65fb5ffc84
+status:
+  HPAReplicas: 6
+  availableReplicas: 5
+  blueGreen: {}
+  canary: {}
+  stableRS: 877894d5b
+  conditions:
+  - lastTransitionTime: "2019-10-25T06:07:29Z"
+    lastUpdateTime: "2019-10-25T06:07:29Z"
+    message: Rollout has minimum availability
+    reason: AvailableReason
+    status: "True"
+    type: Available
+  - lastTransitionTime: "2019-10-28T04:52:55Z"
+    lastUpdateTime: "2019-10-28T04:52:55Z"
+    message: ReplicaSet "canary-demo-65fb5ffc84" has timed out progressing.
+    reason: ProgressDeadlineExceeded
+    status: "False"
+    type: Progressing
+  currentPodHash: 65fb5ffc84
+  currentStepHash: f64cdc9d
+  currentStepIndex: 0
+  observedGeneration: "429"
+  readyReplicas: 5
+  replicas: 6
+  selector: app=canary-demo
+  updatedReplicas: 1

pkg/kubectl-argo-rollouts/info/testdata/canary/canary-rs.yaml

@@ -1,4 +1,4 @@
-apiVersion: extensions/v1beta1
+apiVersion: apps/v1
 kind: ReplicaSet
 metadata:
   annotations:

pkg/kubectl-argo-rollouts/info/testdata/canary/old-rs.yaml

@@ -1,4 +1,4 @@
-apiVersion: extensions/v1beta1
+apiVersion: apps/v1
 kind: ReplicaSet
 metadata:
   annotations:

pkg/kubectl-argo-rollouts/info/testdata/canary/stable-rs.yaml

@@ -1,4 +1,4 @@
-apiVersion: extensions/v1beta1
+apiVersion: apps/v1
 kind: ReplicaSet
 metadata:
   annotations:

rollout/temlateref.go

@@ -34,7 +34,7 @@ type knownKindInfo struct {
 }
 
 var (
-	infoByGroupKind = map[schema.GroupKind]knownKindInfo{
+	InfoByGroupKind = map[schema.GroupKind]knownKindInfo{
 		{Kind: "PodTemplate"}: {
 			TemplatePath: []string{"template"},
 		},
@@ -133,7 +133,7 @@ func (r *informerBasedTemplateResolver) Resolve(rollout *v1alpha1.Rollout) error
 
 	gvk := schema.FromAPIVersionAndKind(rollout.Spec.WorkloadRef.APIVersion, rollout.Spec.WorkloadRef.Kind)
 
-	info, ok := infoByGroupKind[gvk.GroupKind()]
+	info, ok := InfoByGroupKind[gvk.GroupKind()]
 	if !ok {
 		return fmt.Errorf("workload of type %s/%s is not supported", gvk.Group, gvk.Kind)
 	}