Merge pull request #47693 from kow3ns/sts-e2e

Automatic merge from submit-queue (batch tested with PRs 47726, 47693, 46909, 46812)

Additional e2e for StatefulSet Update

**What this PR does / why we need it**:
This PR adds additional e2e tests for StatefulSet update

fixes: #46942

```release-note
NONE
```

pkg/controller/statefulset/stateful_set_control.go

@@ -488,17 +488,9 @@ func (ssc *defaultStatefulSetControl) updateStatefulSet(
 	}
 	// we terminate the Pod with the largest ordinal that does not match the update revision.
 	for target := len(replicas) - 1; target >= updateMin; target-- {
-		// all replicas should be healthy before an update progresses we allow termination of the firstUnhealthy
-		// Pod in any state allow for rolling back a failed update.
-		if !isRunningAndReady(replicas[target]) && replicas[target] != firstUnhealthyPod {
-			glog.V(4).Infof(
-				"StatefulSet %s/%s is waiting for Pod %s to be Running and Ready prior to update",
-				set.Namespace,
-				set.Name,
-				firstUnhealthyPod.Name)
-			return &status, nil
-		}
-		if getPodRevision(replicas[target]) != updateRevision.Name {
+
+		// delete the Pod if it is not already terminating and does not match the update revision.
+		if getPodRevision(replicas[target]) != updateRevision.Name && !isTerminating(replicas[target]) {
 			glog.V(4).Infof("StatefulSet %s/%s terminating Pod %s for update",
 				set.Namespace,
 				set.Name,
@@ -507,6 +499,17 @@ func (ssc *defaultStatefulSetControl) updateStatefulSet(
 			status.CurrentReplicas--
 			return &status, err
 		}
+
+		// wait for unhealthy Pods on update
+		if !isHealthy(replicas[target]) {
+			glog.V(4).Infof(
+				"StatefulSet %s/%s is waiting for Pod %s to update",
+				set.Namespace,
+				set.Name,
+				replicas[target].Name)
+			return &status, nil
+		}
+
 	}
 	return &status, nil
 }

test/e2e/framework/statefulset_utils.go

@@ -19,6 +19,8 @@ package framework
 import (
 	"fmt"
 	"path/filepath"
+	"regexp"
+	"sort"
 	"strconv"
 	"strings"
 	"time"
@@ -33,6 +35,7 @@ import (
 	"k8s.io/apimachinery/pkg/util/sets"
 	"k8s.io/apimachinery/pkg/util/wait"
 	utilyaml "k8s.io/apimachinery/pkg/util/yaml"
+
 	"k8s.io/kubernetes/pkg/api"
 	"k8s.io/kubernetes/pkg/api/v1"
 	podutil "k8s.io/kubernetes/pkg/api/v1/pod"
@@ -96,6 +99,15 @@ func NewStatefulSetTester(c clientset.Interface) *StatefulSetTester {
 	return &StatefulSetTester{c}
 }
 
+// GetStatefulSet gets the StatefulSet named name in namespace.
+func (s *StatefulSetTester) GetStatefulSet(namespace, name string) *apps.StatefulSet {
+	ss, err := s.c.Apps().StatefulSets(namespace).Get(name, metav1.GetOptions{})
+	if err != nil {
+		Failf("Failed to get StatefulSet %s/%s: %v", namespace, name, err)
+	}
+	return ss
+}
+
 // CreateStatefulSet creates a StatefulSet from the manifest at manifestPath in the Namespace ns using kubectl create.
 func (s *StatefulSetTester) CreateStatefulSet(manifestPath, ns string) *apps.StatefulSet {
 	mkpath := func(file string) string {
@@ -324,21 +336,163 @@ func (s *StatefulSetTester) WaitForState(ss *apps.StatefulSet, until func(*apps.
 			return until(ssGet, podList)
 		})
 	if pollErr != nil {
-		Failf("Failed waiting for pods to enter running: %v", pollErr)
+		Failf("Failed waiting for state update: %v", pollErr)
 	}
 }
 
+// WaitForStatus waits for the StatefulSetStatus's ObservedGeneration to be greater than or equal to set's Generation.
+// The returned StatefulSet contains such a StatefulSetStatus
+func (s *StatefulSetTester) WaitForStatus(set *apps.StatefulSet) *apps.StatefulSet {
+	s.WaitForState(set, func(set2 *apps.StatefulSet, pods *v1.PodList) (bool, error) {
+		if set2.Status.ObservedGeneration != nil && *set2.Status.ObservedGeneration >= set.Generation {
+			set = set2
+			return true, nil
+		}
+		return false, nil
+	})
+	return set
+}
+
 // WaitForRunningAndReady waits for numStatefulPods in ss to be Running and Ready.
 func (s *StatefulSetTester) WaitForRunningAndReady(numStatefulPods int32, ss *apps.StatefulSet) {
 	s.waitForRunning(numStatefulPods, ss, true)
 }
 
+// WaitForPodReady waits for the Pod named podName in set to exist and have a Ready condition.
+func (s *StatefulSetTester) WaitForPodReady(set *apps.StatefulSet, podName string) (*apps.StatefulSet, *v1.PodList) {
+	var pods *v1.PodList
+	s.WaitForState(set, func(set2 *apps.StatefulSet, pods2 *v1.PodList) (bool, error) {
+		set = set2
+		pods = pods2
+		for i := range pods.Items {
+			if pods.Items[i].Name == podName {
+				return podutil.IsPodReady(&pods.Items[i]), nil
+			}
+		}
+		return false, nil
+	})
+	return set, pods
+
+}
+
+// WaitForPodNotReady waist for the Pod named podName in set to exist and to not have a Ready condition.
+func (s *StatefulSetTester) WaitForPodNotReady(set *apps.StatefulSet, podName string) (*apps.StatefulSet, *v1.PodList) {
+	var pods *v1.PodList
+	s.WaitForState(set, func(set2 *apps.StatefulSet, pods2 *v1.PodList) (bool, error) {
+		set = set2
+		pods = pods2
+		for i := range pods.Items {
+			if pods.Items[i].Name == podName {
+				return !podutil.IsPodReady(&pods.Items[i]), nil
+			}
+		}
+		return false, nil
+	})
+	return set, pods
+
+}
+
+// WaitForRollingUpdate waits for all Pods in set to exist and have the correct revision and for the RollingUpdate to
+// complete. set must have a RollingUpdateStatefulSetStrategyType.
+func (s *StatefulSetTester) WaitForRollingUpdate(set *apps.StatefulSet) (*apps.StatefulSet, *v1.PodList) {
+	var pods *v1.PodList
+	if set.Spec.UpdateStrategy.Type != apps.RollingUpdateStatefulSetStrategyType {
+		Failf("StatefulSet %s/%s attempt to wait for rolling update with updateStrategy %s",
+			set.Namespace,
+			set.Name,
+			set.Spec.UpdateStrategy.Type)
+	}
+	s.WaitForState(set, func(set2 *apps.StatefulSet, pods2 *v1.PodList) (bool, error) {
+		set = set2
+		pods = pods2
+		if len(pods.Items) < int(*set.Spec.Replicas) {
+			return false, nil
+		}
+		if set.Status.UpdateRevision != set.Status.CurrentRevision {
+			Logf("Waiting for StatefulSet %s/%s to complete update",
+				set.Namespace,
+				set.Name,
+			)
+			s.SortStatefulPods(pods)
+			for i := range pods.Items {
+				if pods.Items[i].Labels[apps.StatefulSetRevisionLabel] != set.Status.UpdateRevision {
+					Logf("Waiting for Pod %s/%s to have revision %s update revision %s",
+						pods.Items[i].Namespace,
+						pods.Items[i].Name,
+						set.Status.UpdateRevision,
+						pods.Items[i].Labels[apps.StatefulSetRevisionLabel])
+				}
+			}
+			return false, nil
+		}
+		return true, nil
+	})
+	return set, pods
+}
+
+// WaitForPartitionedRollingUpdate waits for all Pods in set to exist and have the correct revision. set must have
+// a RollingUpdateStatefulSetStrategyType with a non-nil RollingUpdate and Partition. All Pods with ordinals less
+// than or equal to the Partition are expected to be at set's current revision. All other Pods are expected to be
+// at its update revision.
+func (s *StatefulSetTester) WaitForPartitionedRollingUpdate(set *apps.StatefulSet) (*apps.StatefulSet, *v1.PodList) {
+	var pods *v1.PodList
+	if set.Spec.UpdateStrategy.Type != apps.RollingUpdateStatefulSetStrategyType {
+		Failf("StatefulSet %s/%s attempt to wait for partitioned update with updateStrategy %s",
+			set.Namespace,
+			set.Name,
+			set.Spec.UpdateStrategy.Type)
+	}
+	if set.Spec.UpdateStrategy.RollingUpdate == nil || set.Spec.UpdateStrategy.RollingUpdate.Partition == nil {
+		Failf("StatefulSet %s/%s attempt to wait for partitioned update with nil RollingUpdate or nil Partition",
+			set.Namespace,
+			set.Name)
+	}
+	s.WaitForState(set, func(set2 *apps.StatefulSet, pods2 *v1.PodList) (bool, error) {
+		set = set2
+		pods = pods2
+		partition := int(*set.Spec.UpdateStrategy.RollingUpdate.Partition)
+		if len(pods.Items) < int(*set.Spec.Replicas) {
+			return false, nil
+		}
+		if partition <= 0 && set.Status.UpdateRevision != set.Status.CurrentRevision {
+			Logf("Waiting for StatefulSet %s/%s to complete update",
+				set.Namespace,
+				set.Name,
+			)
+			s.SortStatefulPods(pods)
+			for i := range pods.Items {
+				if pods.Items[i].Labels[apps.StatefulSetRevisionLabel] != set.Status.UpdateRevision {
+					Logf("Waiting for Pod %s/%s to have revision %s update revision %s",
+						pods.Items[i].Namespace,
+						pods.Items[i].Name,
+						set.Status.UpdateRevision,
+						pods.Items[i].Labels[apps.StatefulSetRevisionLabel])
+				}
+			}
+			return false, nil
+		} else {
+			for i := int(*set.Spec.Replicas) - 1; i >= partition; i-- {
+				if pods.Items[i].Labels[apps.StatefulSetRevisionLabel] != set.Status.UpdateRevision {
+					Logf("Waiting for Pod %s/%s to have revision %s update revision %s",
+						pods.Items[i].Namespace,
+						pods.Items[i].Name,
+						set.Status.UpdateRevision,
+						pods.Items[i].Labels[apps.StatefulSetRevisionLabel])
+					return false, nil
+				}
+			}
+		}
+		return true, nil
+	})
+	return set, pods
+}
+
 // WaitForRunningAndReady waits for numStatefulPods in ss to be Running and not Ready.
 func (s *StatefulSetTester) WaitForRunningAndNotReady(numStatefulPods int32, ss *apps.StatefulSet) {
 	s.waitForRunning(numStatefulPods, ss, false)
 }
 
-// BreakProbe breaks the readiness probe for Nginx StatefulSet containers.
+// BreakProbe breaks the readiness probe for Nginx StatefulSet containers in ss.
 func (s *StatefulSetTester) BreakProbe(ss *apps.StatefulSet, probe *v1.Probe) error {
 	path := probe.HTTPGet.Path
 	if path == "" {
@@ -348,7 +502,19 @@ func (s *StatefulSetTester) BreakProbe(ss *apps.StatefulSet, probe *v1.Probe) er
 	return s.ExecInStatefulPods(ss, cmd)
 }
 
-// RestoreProbe restores the readiness probe for Nginx StatefulSet containers.
+// BreakProbe breaks the readiness probe for Nginx StatefulSet containers in pod.
+func (s *StatefulSetTester) BreakPodProbe(ss *apps.StatefulSet, pod *v1.Pod, probe *v1.Probe) error {
+	path := probe.HTTPGet.Path
+	if path == "" {
+		return fmt.Errorf("Path expected to be not empty: %v", path)
+	}
+	cmd := fmt.Sprintf("mv -v /usr/share/nginx/html%v /tmp/", path)
+	stdout, err := RunHostCmd(pod.Namespace, pod.Name, cmd)
+	Logf("stdout of %v on %v: %v", cmd, pod.Name, stdout)
+	return err
+}
+
+// RestoreProbe restores the readiness probe for Nginx StatefulSet containers in ss.
 func (s *StatefulSetTester) RestoreProbe(ss *apps.StatefulSet, probe *v1.Probe) error {
 	path := probe.HTTPGet.Path
 	if path == "" {
@@ -358,6 +524,18 @@ func (s *StatefulSetTester) RestoreProbe(ss *apps.StatefulSet, probe *v1.Probe)
 	return s.ExecInStatefulPods(ss, cmd)
 }
 
+// RestoreProbe restores the readiness probe for Nginx StatefulSet containers in pod.
+func (s *StatefulSetTester) RestorePodProbe(ss *apps.StatefulSet, pod *v1.Pod, probe *v1.Probe) error {
+	path := probe.HTTPGet.Path
+	if path == "" {
+		return fmt.Errorf("Path expected to be not empty: %v", path)
+	}
+	cmd := fmt.Sprintf("mv -v /tmp%v /usr/share/nginx/html/", path)
+	stdout, err := RunHostCmd(pod.Namespace, pod.Name, cmd)
+	Logf("stdout of %v on %v: %v", cmd, pod.Name, stdout)
+	return err
+}
+
 // SetHealthy updates the StatefulSet InitAnnotation to true in order to set a StatefulSet Pod to be Running and Ready.
 func (s *StatefulSetTester) SetHealthy(ss *apps.StatefulSet) {
 	podList := s.GetPodList(ss)
@@ -443,6 +621,11 @@ func (p *StatefulSetTester) CheckServiceName(ss *apps.StatefulSet, expectedServi
 	return nil
 }
 
+// SortStatefulPods sorts pods by their ordinals
+func (s *StatefulSetTester) SortStatefulPods(pods *v1.PodList) {
+	sort.Sort(statefulPodsByOrdinal(pods.Items))
+}
+
 // DeleteAllStatefulSets deletes all StatefulSet API Objects in Namespace ns.
 func DeleteAllStatefulSets(c clientset.Interface, ns string) {
 	sst := &StatefulSetTester{c: c}
@@ -613,3 +796,31 @@ func NewStatefulSet(name, ns, governingSvcName string, replicas int32, statefulP
 func SetStatefulSetInitializedAnnotation(ss *apps.StatefulSet, value string) {
 	ss.Spec.Template.ObjectMeta.Annotations["pod.alpha.kubernetes.io/initialized"] = value
 }
+
+var statefulPodRegex = regexp.MustCompile("(.*)-([0-9]+)$")
+
+func getStatefulPodOrdinal(pod *v1.Pod) int {
+	ordinal := -1
+	subMatches := statefulPodRegex.FindStringSubmatch(pod.Name)
+	if len(subMatches) < 3 {
+		return ordinal
+	}
+	if i, err := strconv.ParseInt(subMatches[2], 10, 32); err == nil {
+		ordinal = int(i)
+	}
+	return ordinal
+}
+
+type statefulPodsByOrdinal []v1.Pod
+
+func (sp statefulPodsByOrdinal) Len() int {
+	return len(sp)
+}
+
+func (sp statefulPodsByOrdinal) Swap(i, j int) {
+	sp[i], sp[j] = sp[j], sp[i]
+}
+
+func (sp statefulPodsByOrdinal) Less(i, j int) bool {
+	return getStatefulPodOrdinal(&sp[i]) < getStatefulPodOrdinal(&sp[j])
+}

test/e2e/statefulset.go

@@ -247,44 +247,417 @@ var _ = framework.KubeDescribe("StatefulSet", func() {
 			sst.Saturate(ss)
 		})
 
-		It("should allow template updates", func() {
-			By("Creating stateful set " + ssName + " in namespace " + ns)
+		It("should perform rolling updates and roll backs of template modifications", func() {
+			By("Creating a new StatefulSet")
 			testProbe := &v1.Probe{Handler: v1.Handler{HTTPGet: &v1.HTTPGetAction{
 				Path: "/index.html",
 				Port: intstr.IntOrString{IntVal: 80}}}}
-			ss := framework.NewStatefulSet("ss2", ns, headlessSvcName, 2, nil, nil, labels)
+			ss := framework.NewStatefulSet("ss2", ns, headlessSvcName, 3, nil, nil, labels)
 			ss.Spec.Template.Spec.Containers[0].ReadinessProbe = testProbe
 			ss, err := c.Apps().StatefulSets(ns).Create(ss)
 			Expect(err).NotTo(HaveOccurred())
-
 			sst := framework.NewStatefulSetTester(c)
-
 			sst.WaitForRunningAndReady(*ss.Spec.Replicas, ss)
-
+			ss = sst.WaitForStatus(ss)
+			currentRevision, updateRevision := ss.Status.CurrentRevision, ss.Status.UpdateRevision
+			Expect(currentRevision).To(Equal(updateRevision),
+				fmt.Sprintf("StatefulSet %s/%s created with update revision %s not equal to current revision %s",
+					ss.Namespace, ss.Name, updateRevision, currentRevision))
+			pods := sst.GetPodList(ss)
+			for i := range pods.Items {
+				Expect(pods.Items[i].Labels[apps.StatefulSetRevisionLabel]).To(Equal(currentRevision),
+					fmt.Sprintf("Pod %s/%s revision %s is not equal to current revision %s",
+						pods.Items[i].Namespace,
+						pods.Items[i].Name,
+						pods.Items[i].Labels[apps.StatefulSetRevisionLabel],
+						currentRevision))
+			}
+			sst.SortStatefulPods(pods)
+			sst.BreakPodProbe(ss, &pods.Items[1], testProbe)
+			Expect(err).NotTo(HaveOccurred())
+			ss, pods = sst.WaitForPodNotReady(ss, pods.Items[1].Name)
 			newImage := newNginxImage
 			oldImage := ss.Spec.Template.Spec.Containers[0].Image
-			By(fmt.Sprintf("Updating stateful set template: update image from %s to %s", oldImage, newImage))
+
+			By(fmt.Sprintf("Updating StatefulSet template: update image from %s to %s", oldImage, newImage))
 			Expect(oldImage).NotTo(Equal(newImage), "Incorrect test setup: should update to a different image")
-			_, err = framework.UpdateStatefulSetWithRetries(c, ns, ss.Name, func(update *apps.StatefulSet) {
+			ss, err = framework.UpdateStatefulSetWithRetries(c, ns, ss.Name, func(update *apps.StatefulSet) {
 				update.Spec.Template.Spec.Containers[0].Image = newImage
 			})
 			Expect(err).NotTo(HaveOccurred())
 
-			sst.WaitForState(ss, func(set *apps.StatefulSet, pods *v1.PodList) (bool, error) {
-				if len(pods.Items) < 2 {
-					return false, nil
+			By("Creating a new revision")
+			ss = sst.WaitForStatus(ss)
+			currentRevision, updateRevision = ss.Status.CurrentRevision, ss.Status.UpdateRevision
+			Expect(currentRevision).NotTo(Equal(updateRevision),
+				"Current revision should not equal update revision during rolling update")
+
+			By("Updating Pods in reverse ordinal order")
+			pods = sst.GetPodList(ss)
+			sst.SortStatefulPods(pods)
+			sst.RestorePodProbe(ss, &pods.Items[1], testProbe)
+			ss, pods = sst.WaitForPodReady(ss, pods.Items[1].Name)
+			ss, pods = sst.WaitForRollingUpdate(ss)
+			Expect(ss.Status.CurrentRevision).To(Equal(updateRevision),
+				fmt.Sprintf("StatefulSet %s/%s current revision %s does not equal updste revision %s on update completion",
+					ss.Namespace,
+					ss.Name,
+					ss.Status.CurrentRevision,
+					updateRevision))
+			for i := range pods.Items {
+				Expect(pods.Items[i].Spec.Containers[0].Image).To(Equal(newImage),
+					fmt.Sprintf(" Pod %s/%s has image %s not have new image %s",
+						pods.Items[i].Namespace,
+						pods.Items[i].Name,
+						pods.Items[i].Spec.Containers[0].Image,
+						newImage))
+				Expect(pods.Items[i].Labels[apps.StatefulSetRevisionLabel]).To(Equal(updateRevision),
+					fmt.Sprintf("Pod %s/%s revision %s is not equal to update revision %s",
+						pods.Items[i].Namespace,
+						pods.Items[i].Name,
+						pods.Items[i].Labels[apps.StatefulSetRevisionLabel],
+						updateRevision))
+			}
+
+			By("Rolling back to a previous revision")
+			sst.BreakPodProbe(ss, &pods.Items[1], testProbe)
+			Expect(err).NotTo(HaveOccurred())
+			ss, pods = sst.WaitForPodNotReady(ss, pods.Items[1].Name)
+			priorRevision := currentRevision
+			currentRevision, updateRevision = ss.Status.CurrentRevision, ss.Status.UpdateRevision
+			ss, err = framework.UpdateStatefulSetWithRetries(c, ns, ss.Name, func(update *apps.StatefulSet) {
+				update.Spec.Template.Spec.Containers[0].Image = oldImage
+			})
+			Expect(err).NotTo(HaveOccurred())
+			ss = sst.WaitForStatus(ss)
+			currentRevision, updateRevision = ss.Status.CurrentRevision, ss.Status.UpdateRevision
+			Expect(currentRevision).NotTo(Equal(updateRevision),
+				"Current revision should not equal update revision during roll bakc")
+			Expect(priorRevision).To(Equal(updateRevision),
+				"Prior revision should equal update revision during roll back")
+
+			By("Rolling back update in reverse ordinal order")
+			pods = sst.GetPodList(ss)
+			sst.SortStatefulPods(pods)
+			sst.RestorePodProbe(ss, &pods.Items[1], testProbe)
+			ss, pods = sst.WaitForPodReady(ss, pods.Items[1].Name)
+			ss, pods = sst.WaitForRollingUpdate(ss)
+			Expect(ss.Status.CurrentRevision).To(Equal(priorRevision),
+				fmt.Sprintf("StatefulSet %s/%s current revision %s does not equal prior revision %s on rollback completion",
+					ss.Namespace,
+					ss.Name,
+					ss.Status.CurrentRevision,
+					updateRevision))
+
+			for i := range pods.Items {
+				Expect(pods.Items[i].Spec.Containers[0].Image).To(Equal(oldImage),
+					fmt.Sprintf("Pod %s/%s has image %s not equal to previous image %s",
+						pods.Items[i].Namespace,
+						pods.Items[i].Name,
+						pods.Items[i].Spec.Containers[0].Image,
+						oldImage))
+				Expect(pods.Items[i].Labels[apps.StatefulSetRevisionLabel]).To(Equal(priorRevision),
+					fmt.Sprintf("Pod %s/%s revision %s is not equal to prior revision %s",
+						pods.Items[i].Namespace,
+						pods.Items[i].Name,
+						pods.Items[i].Labels[apps.StatefulSetRevisionLabel],
+						priorRevision))
+			}
+		})
+
+		It("should perform canary updates and phased rolling updates of template modifications", func() {
+			By("Creating a new StaefulSet")
+			testProbe := &v1.Probe{Handler: v1.Handler{HTTPGet: &v1.HTTPGetAction{
+				Path: "/index.html",
+				Port: intstr.IntOrString{IntVal: 80}}}}
+			ss := framework.NewStatefulSet("ss2", ns, headlessSvcName, 3, nil, nil, labels)
+			ss.Spec.Template.Spec.Containers[0].ReadinessProbe = testProbe
+			ss.Spec.UpdateStrategy = apps.StatefulSetUpdateStrategy{
+				Type: apps.RollingUpdateStatefulSetStrategyType,
+				RollingUpdate: func() *apps.RollingUpdateStatefulSetStrategy {
+					return &apps.RollingUpdateStatefulSetStrategy{
+						Partition: func() *int32 {
+							i := int32(3)
+							return &i
+						}()}
+				}(),
+			}
+			ss, err := c.Apps().StatefulSets(ns).Create(ss)
+			Expect(err).NotTo(HaveOccurred())
+			sst := framework.NewStatefulSetTester(c)
+			sst.WaitForRunningAndReady(*ss.Spec.Replicas, ss)
+			ss = sst.WaitForStatus(ss)
+			currentRevision, updateRevision := ss.Status.CurrentRevision, ss.Status.UpdateRevision
+			Expect(currentRevision).To(Equal(updateRevision),
+				fmt.Sprintf("StatefulSet %s/%s created with update revision %s not equal to current revision %s",
+					ss.Namespace, ss.Name, updateRevision, currentRevision))
+			pods := sst.GetPodList(ss)
+			for i := range pods.Items {
+				Expect(pods.Items[i].Labels[apps.StatefulSetRevisionLabel]).To(Equal(currentRevision),
+					fmt.Sprintf("Pod %s/%s revision %s is not equal to currentRevision %s",
+						pods.Items[i].Namespace,
+						pods.Items[i].Name,
+						pods.Items[i].Labels[apps.StatefulSetRevisionLabel],
+						currentRevision))
+			}
+			newImage := newNginxImage
+			oldImage := ss.Spec.Template.Spec.Containers[0].Image
+
+			By(fmt.Sprintf("Updating stateful set template: update image from %s to %s", oldImage, newImage))
+			Expect(oldImage).NotTo(Equal(newImage), "Incorrect test setup: should update to a different image")
+			ss, err = framework.UpdateStatefulSetWithRetries(c, ns, ss.Name, func(update *apps.StatefulSet) {
+				update.Spec.Template.Spec.Containers[0].Image = newImage
+			})
+			Expect(err).NotTo(HaveOccurred())
+
+			By("Creating a new revision")
+			ss = sst.WaitForStatus(ss)
+			currentRevision, updateRevision = ss.Status.CurrentRevision, ss.Status.UpdateRevision
+			Expect(currentRevision).NotTo(Equal(updateRevision),
+				"Current revision should not equal update revision during rolling update")
+
+			By("Not applying an update when the partition is greater than the number of replicas")
+			for i := range pods.Items {
+				Expect(pods.Items[i].Spec.Containers[0].Image).To(Equal(oldImage),
+					fmt.Sprintf("Pod %s/%s has image %s not equal to current image %s",
+						pods.Items[i].Namespace,
+						pods.Items[i].Name,
+						pods.Items[i].Spec.Containers[0].Image,
+						oldImage))
+				Expect(pods.Items[i].Labels[apps.StatefulSetRevisionLabel]).To(Equal(currentRevision),
+					fmt.Sprintf("Pod %s/%s has revision %s not equal to current revision %s",
+						pods.Items[i].Namespace,
+						pods.Items[i].Name,
+						pods.Items[i].Labels[apps.StatefulSetRevisionLabel],
+						currentRevision))
+			}
+
+			By("By performing a canary update")
+			ss.Spec.UpdateStrategy = apps.StatefulSetUpdateStrategy{
+				Type: apps.RollingUpdateStatefulSetStrategyType,
+				RollingUpdate: func() *apps.RollingUpdateStatefulSetStrategy {
+					return &apps.RollingUpdateStatefulSetStrategy{
+						Partition: func() *int32 {
+							i := int32(2)
+							return &i
+						}()}
+				}(),
+			}
+			ss, err = framework.UpdateStatefulSetWithRetries(c, ns, ss.Name, func(update *apps.StatefulSet) {
+				update.Spec.UpdateStrategy = apps.StatefulSetUpdateStrategy{
+					Type: apps.RollingUpdateStatefulSetStrategyType,
+					RollingUpdate: func() *apps.RollingUpdateStatefulSetStrategy {
+						return &apps.RollingUpdateStatefulSetStrategy{
+							Partition: func() *int32 {
+								i := int32(2)
+								return &i
+							}()}
+					}(),
 				}
-				for i := range pods.Items {
-					if pods.Items[i].Spec.Containers[0].Image != newImage {
-						framework.Logf("Waiting for pod %s to have image %s current image %s",
+			})
+			Expect(err).NotTo(HaveOccurred())
+			ss, pods = sst.WaitForPartitionedRollingUpdate(ss)
+			for i := range pods.Items {
+				if i < int(*ss.Spec.UpdateStrategy.RollingUpdate.Partition) {
+					Expect(pods.Items[i].Spec.Containers[0].Image).To(Equal(oldImage),
+						fmt.Sprintf("Pod %s/%s has image %s not equal to current image %s",
+							pods.Items[i].Namespace,
 							pods.Items[i].Name,
-							newImage,
-							pods.Items[i].Spec.Containers[0].Image)
-						return false, nil
+							pods.Items[i].Spec.Containers[0].Image,
+							oldImage))
+					Expect(pods.Items[i].Labels[apps.StatefulSetRevisionLabel]).To(Equal(currentRevision),
+						fmt.Sprintf("Pod %s/%s has revision %s not equal to current revision %s",
+							pods.Items[i].Namespace,
+							pods.Items[i].Name,
+							pods.Items[i].Labels[apps.StatefulSetRevisionLabel],
+							currentRevision))
+				} else {
+					Expect(pods.Items[i].Spec.Containers[0].Image).To(Equal(newImage),
+						fmt.Sprintf("Pod %s/%s has image %s not equal to new image  %s",
+							pods.Items[i].Namespace,
+							pods.Items[i].Name,
+							pods.Items[i].Spec.Containers[0].Image,
+							newImage))
+					Expect(pods.Items[i].Labels[apps.StatefulSetRevisionLabel]).To(Equal(updateRevision),
+						fmt.Sprintf("Pod %s/%s has revision %s not equal to new revision %s",
+							pods.Items[i].Namespace,
+							pods.Items[i].Name,
+							pods.Items[i].Labels[apps.StatefulSetRevisionLabel],
+							updateRevision))
+				}
+			}
+
+			By("Restoring Pods to the correct revision when they are deleted")
+			sst.DeleteStatefulPodAtIndex(0, ss)
+			sst.DeleteStatefulPodAtIndex(2, ss)
+			sst.WaitForRunningAndReady(3, ss)
+			ss = sst.GetStatefulSet(ss.Namespace, ss.Name)
+			pods = sst.GetPodList(ss)
+			for i := range pods.Items {
+				if i < int(*ss.Spec.UpdateStrategy.RollingUpdate.Partition) {
+					Expect(pods.Items[i].Spec.Containers[0].Image).To(Equal(oldImage),
+						fmt.Sprintf("Pod %s/%s has image %s not equal to current image %s",
+							pods.Items[i].Namespace,
+							pods.Items[i].Name,
+							pods.Items[i].Spec.Containers[0].Image,
+							oldImage))
+					Expect(pods.Items[i].Labels[apps.StatefulSetRevisionLabel]).To(Equal(currentRevision),
+						fmt.Sprintf("Pod %s/%s has revision %s not equal to current revision %s",
+							pods.Items[i].Namespace,
+							pods.Items[i].Name,
+							pods.Items[i].Labels[apps.StatefulSetRevisionLabel],
+							currentRevision))
+				} else {
+					Expect(pods.Items[i].Spec.Containers[0].Image).To(Equal(newImage),
+						fmt.Sprintf("Pod %s/%s has image %s not equal to new image  %s",
+							pods.Items[i].Namespace,
+							pods.Items[i].Name,
+							pods.Items[i].Spec.Containers[0].Image,
+							newImage))
+					Expect(pods.Items[i].Labels[apps.StatefulSetRevisionLabel]).To(Equal(updateRevision),
+						fmt.Sprintf("Pod %s/%s has revision %s not equal to new revision %s",
+							pods.Items[i].Namespace,
+							pods.Items[i].Name,
+							pods.Items[i].Labels[apps.StatefulSetRevisionLabel],
+							updateRevision))
+				}
+			}
+
+			By("Performing a phased rolling update")
+			for i := int(*ss.Spec.UpdateStrategy.RollingUpdate.Partition) - 1; i >= 0; i-- {
+				ss, err = framework.UpdateStatefulSetWithRetries(c, ns, ss.Name, func(update *apps.StatefulSet) {
+					update.Spec.UpdateStrategy = apps.StatefulSetUpdateStrategy{
+						Type: apps.RollingUpdateStatefulSetStrategyType,
+						RollingUpdate: func() *apps.RollingUpdateStatefulSetStrategy {
+							j := int32(i)
+							return &apps.RollingUpdateStatefulSetStrategy{
+								Partition: &j,
+							}
+						}(),
+					}
+				})
+				Expect(err).NotTo(HaveOccurred())
+				ss, pods = sst.WaitForPartitionedRollingUpdate(ss)
+				for i := range pods.Items {
+					if i < int(*ss.Spec.UpdateStrategy.RollingUpdate.Partition) {
+						Expect(pods.Items[i].Spec.Containers[0].Image).To(Equal(oldImage),
+							fmt.Sprintf("Pod %s/%s has image %s not equal to current image %s",
+								pods.Items[i].Namespace,
+								pods.Items[i].Name,
+								pods.Items[i].Spec.Containers[0].Image,
+								oldImage))
+						Expect(pods.Items[i].Labels[apps.StatefulSetRevisionLabel]).To(Equal(currentRevision),
+							fmt.Sprintf("Pod %s/%s has revision %s not equal to current revision %s",
+								pods.Items[i].Namespace,
+								pods.Items[i].Name,
+								pods.Items[i].Labels[apps.StatefulSetRevisionLabel],
+								currentRevision))
+					} else {
+						Expect(pods.Items[i].Spec.Containers[0].Image).To(Equal(newImage),
+							fmt.Sprintf("Pod %s/%s has image %s not equal to new image  %s",
+								pods.Items[i].Namespace,
+								pods.Items[i].Name,
+								pods.Items[i].Spec.Containers[0].Image,
+								newImage))
+						Expect(pods.Items[i].Labels[apps.StatefulSetRevisionLabel]).To(Equal(updateRevision),
+							fmt.Sprintf("Pod %s/%s has revision %s not equal to new revision %s",
+								pods.Items[i].Namespace,
+								pods.Items[i].Name,
+								pods.Items[i].Labels[apps.StatefulSetRevisionLabel],
+								updateRevision))
 					}
 				}
-				return true, nil
+			}
+			Expect(ss.Status.CurrentRevision).To(Equal(updateRevision),
+				fmt.Sprintf("StatefulSet %s/%s current revision %s does not equal update revison %s on update completion",
+					ss.Namespace,
+					ss.Name,
+					ss.Status.CurrentRevision,
+					updateRevision))
+
+		})
+
+		It("should implement legacy replacement when the update strategy is OnDelete", func() {
+			By("Creating a new StatefulSet")
+			testProbe := &v1.Probe{Handler: v1.Handler{HTTPGet: &v1.HTTPGetAction{
+				Path: "/index.html",
+				Port: intstr.IntOrString{IntVal: 80}}}}
+			ss := framework.NewStatefulSet("ss2", ns, headlessSvcName, 3, nil, nil, labels)
+			ss.Spec.Template.Spec.Containers[0].ReadinessProbe = testProbe
+			ss.Spec.UpdateStrategy = apps.StatefulSetUpdateStrategy{
+				Type: apps.OnDeleteStatefulSetStrategyType,
+			}
+			ss, err := c.Apps().StatefulSets(ns).Create(ss)
+			Expect(err).NotTo(HaveOccurred())
+			sst := framework.NewStatefulSetTester(c)
+			sst.WaitForRunningAndReady(*ss.Spec.Replicas, ss)
+			ss = sst.WaitForStatus(ss)
+			currentRevision, updateRevision := ss.Status.CurrentRevision, ss.Status.UpdateRevision
+			Expect(currentRevision).To(Equal(updateRevision),
+				fmt.Sprintf("StatefulSet %s/%s created with update revision %s not equal to current revision %s",
+					ss.Namespace, ss.Name, updateRevision, currentRevision))
+			pods := sst.GetPodList(ss)
+			for i := range pods.Items {
+				Expect(pods.Items[i].Labels[apps.StatefulSetRevisionLabel]).To(Equal(currentRevision),
+					fmt.Sprintf("Pod %s/%s revision %s is not equal to current revision %s",
+						pods.Items[i].Namespace,
+						pods.Items[i].Name,
+						pods.Items[i].Labels[apps.StatefulSetRevisionLabel],
+						currentRevision))
+			}
+
+			By("Restoring Pods to the current revision")
+			sst.DeleteStatefulPodAtIndex(0, ss)
+			sst.DeleteStatefulPodAtIndex(1, ss)
+			sst.DeleteStatefulPodAtIndex(2, ss)
+			sst.WaitForRunningAndReady(3, ss)
+			ss = sst.GetStatefulSet(ss.Namespace, ss.Name)
+			pods = sst.GetPodList(ss)
+			for i := range pods.Items {
+				Expect(pods.Items[i].Labels[apps.StatefulSetRevisionLabel]).To(Equal(currentRevision),
+					fmt.Sprintf("Pod %s/%s revision %s is not equal to current revision %s",
+						pods.Items[i].Namespace,
+						pods.Items[i].Name,
+						pods.Items[i].Labels[apps.StatefulSetRevisionLabel],
+						currentRevision))
+			}
+			newImage := newNginxImage
+			oldImage := ss.Spec.Template.Spec.Containers[0].Image
+
+			By(fmt.Sprintf("Updating stateful set template: update image from %s to %s", oldImage, newImage))
+			Expect(oldImage).NotTo(Equal(newImage), "Incorrect test setup: should update to a different image")
+			ss, err = framework.UpdateStatefulSetWithRetries(c, ns, ss.Name, func(update *apps.StatefulSet) {
+				update.Spec.Template.Spec.Containers[0].Image = newImage
 			})
+			Expect(err).NotTo(HaveOccurred())
+
+			By("Creating a new revision")
+			ss = sst.WaitForStatus(ss)
+			currentRevision, updateRevision = ss.Status.CurrentRevision, ss.Status.UpdateRevision
+			Expect(currentRevision).NotTo(Equal(updateRevision),
+				"Current revision should not equal update revision during rolling update")
+
+			By("Recreating Pods at the new revision")
+			sst.DeleteStatefulPodAtIndex(0, ss)
+			sst.DeleteStatefulPodAtIndex(1, ss)
+			sst.DeleteStatefulPodAtIndex(2, ss)
+			sst.WaitForRunningAndReady(3, ss)
+			ss = sst.GetStatefulSet(ss.Namespace, ss.Name)
+			pods = sst.GetPodList(ss)
+			for i := range pods.Items {
+				Expect(pods.Items[i].Spec.Containers[0].Image).To(Equal(newImage),
+					fmt.Sprintf("Pod %s/%s has image %s not equal to new image %s",
+						pods.Items[i].Namespace,
+						pods.Items[i].Name,
+						pods.Items[i].Spec.Containers[0].Image,
+						newImage))
+				Expect(pods.Items[i].Labels[apps.StatefulSetRevisionLabel]).To(Equal(updateRevision),
+					fmt.Sprintf("Pod %s/%s has revision %s not equal to current revision %s",
+						pods.Items[i].Namespace,
+						pods.Items[i].Name,
+						pods.Items[i].Labels[apps.StatefulSetRevisionLabel],
+						updateRevision))
+			}
 		})
 
 		It("Scaling should happen in predictable order and halt if any stateful pod is unhealthy", func() {