From 594e228a17d340c14a46b6673b25f4941acc1baf Mon Sep 17 00:00:00 2001
From: Faraaz Khan <faraaz@rationalizeit.us>
Date: Wed, 25 Apr 2018 18:04:25 -0500
Subject: [PATCH] set updated replicas correctly in scale up, scale down
 scenarios as well

---
 .../statefulset/stateful_set_control.go       |  9 ++--
 .../statefulset/stateful_set_control_test.go  | 52 +++++++++++++------
 2 files changed, 42 insertions(+), 19 deletions(-)

diff --git a/pkg/controller/statefulset/stateful_set_control.go b/pkg/controller/statefulset/stateful_set_control.go
index edfc3c5bba5..9634b3d8d6a 100644
--- a/pkg/controller/statefulset/stateful_set_control.go
+++ b/pkg/controller/statefulset/stateful_set_control.go
@@ -380,7 +380,8 @@ func (ssc *defaultStatefulSetControl) updateStatefulSet(
 			}
 			if getPodRevision(replicas[i]) == currentRevision.Name {
 				status.CurrentReplicas--
-			} else if getPodRevision(replicas[i]) == updateRevision.Name {
+			}
+			if getPodRevision(replicas[i]) == updateRevision.Name {
 				status.UpdatedReplicas--
 			}
 			status.Replicas--
@@ -399,7 +400,8 @@ func (ssc *defaultStatefulSetControl) updateStatefulSet(
 			status.Replicas++
 			if getPodRevision(replicas[i]) == currentRevision.Name {
 				status.CurrentReplicas++
-			} else if getPodRevision(replicas[i]) == updateRevision.Name {
+			}
+			if getPodRevision(replicas[i]) == updateRevision.Name {
 				status.UpdatedReplicas++
 			}
 
@@ -480,7 +482,8 @@ func (ssc *defaultStatefulSetControl) updateStatefulSet(
 		}
 		if getPodRevision(condemned[target]) == currentRevision.Name {
 			status.CurrentReplicas--
-		} else if getPodRevision(condemned[target]) == updateRevision.Name {
+		}
+		if getPodRevision(condemned[target]) == updateRevision.Name {
 			status.UpdatedReplicas--
 		}
 		if monotonic {
diff --git a/pkg/controller/statefulset/stateful_set_control_test.go b/pkg/controller/statefulset/stateful_set_control_test.go
index e0de54e4288..6af6630c378 100644
--- a/pkg/controller/statefulset/stateful_set_control_test.go
+++ b/pkg/controller/statefulset/stateful_set_control_test.go
@@ -81,7 +81,6 @@ func TestStatefulSetControl(t *testing.T) {
 		obj func() *apps.StatefulSet
 	}{
 		{CreatesPods, simpleSetFn},
-		{SetsStatusFields, simpleSetFn},
 		{ScalesUp, simpleSetFn},
 		{ScalesDown, simpleSetFn},
 		{ReplacesPods, largeSetFn},
@@ -129,21 +128,6 @@ func CreatesPods(t *testing.T, set *apps.StatefulSet, invariants invariantFunc)
 	if set.Status.Replicas != 3 {
 		t.Error("Failed to scale statefulset to 3 replicas")
 	}
-}
-
-func SetsStatusFields(t *testing.T, set *apps.StatefulSet, invariants invariantFunc) {
-	client := fake.NewSimpleClientset(set)
-	spc, _, ssc, stop := setupController(client)
-	defer close(stop)
-
-	if err := scaleUpStatefulSetControl(set, ssc, spc, invariants); err != nil {
-		t.Errorf("Failed to turn up StatefulSet : %s", err)
-	}
-	var err error
-	set, err = spc.setsLister.StatefulSets(set.Namespace).Get(set.Name)
-	if err != nil {
-		t.Fatalf("Error getting updated StatefulSet: %v", err)
-	}
 	if set.Status.ReadyReplicas != 3 {
 		t.Error("Failed to set ReadyReplicas correctly")
 	}
@@ -172,6 +156,12 @@ func ScalesUp(t *testing.T, set *apps.StatefulSet, invariants invariantFunc) {
 	if set.Status.Replicas != 4 {
 		t.Error("Failed to scale statefulset to 4 replicas")
 	}
+	if set.Status.ReadyReplicas != 4 {
+		t.Error("Failed to set readyReplicas correctly")
+	}
+	if set.Status.UpdatedReplicas != 4 {
+		t.Error("Failed to set updatedReplicas correctly")
+	}
 }
 
 func ScalesDown(t *testing.T, set *apps.StatefulSet, invariants invariantFunc) {
@@ -189,6 +179,12 @@ func ScalesDown(t *testing.T, set *apps.StatefulSet, invariants invariantFunc) {
 	if set.Status.Replicas != 0 {
 		t.Error("Failed to scale statefulset to 0 replicas")
 	}
+	if set.Status.ReadyReplicas != 0 {
+		t.Error("Failed to set readyReplicas correctly")
+	}
+	if set.Status.UpdatedReplicas != 0 {
+		t.Error("Failed to set updatedReplicas correctly")
+	}
 }
 
 func ReplacesPods(t *testing.T, set *apps.StatefulSet, invariants invariantFunc) {
@@ -320,6 +316,12 @@ func CreatePodFailure(t *testing.T, set *apps.StatefulSet, invariants invariantF
 	if set.Status.Replicas != 3 {
 		t.Error("Failed to scale StatefulSet to 3 replicas")
 	}
+	if set.Status.ReadyReplicas != 3 {
+		t.Error("Failed to set readyReplicas correctly")
+	}
+	if set.Status.UpdatedReplicas != 3 {
+		t.Error("Failed to updatedReplicas correctly")
+	}
 }
 
 func UpdatePodFailure(t *testing.T, set *apps.StatefulSet, invariants invariantFunc) {
@@ -340,6 +342,12 @@ func UpdatePodFailure(t *testing.T, set *apps.StatefulSet, invariants invariantF
 	if set.Status.Replicas != 3 {
 		t.Error("Failed to scale StatefulSet to 3 replicas")
 	}
+	if set.Status.ReadyReplicas != 3 {
+		t.Error("Failed to set readyReplicas correctly")
+	}
+	if set.Status.UpdatedReplicas != 3 {
+		t.Error("Failed to set updatedReplicas correctly")
+	}
 
 	// now mutate a pod's identity
 	pods, err := spc.podsLister.List(labels.Everything())
@@ -379,6 +387,12 @@ func UpdateSetStatusFailure(t *testing.T, set *apps.StatefulSet, invariants inva
 	if set.Status.Replicas != 3 {
 		t.Error("Failed to scale StatefulSet to 3 replicas")
 	}
+	if set.Status.ReadyReplicas != 3 {
+		t.Error("Failed to set readyReplicas to 3")
+	}
+	if set.Status.UpdatedReplicas != 3 {
+		t.Error("Failed to set updatedReplicas to 3")
+	}
 }
 
 func PodRecreateDeleteFailure(t *testing.T, set *apps.StatefulSet, invariants invariantFunc) {
@@ -458,6 +472,12 @@ func TestStatefulSetControlScaleDownDeleteError(t *testing.T) {
 	if set.Status.Replicas != 0 {
 		t.Error("Failed to scale statefulset to 0 replicas")
 	}
+	if set.Status.ReadyReplicas != 0 {
+		t.Error("Failed to set readyReplicas to 0")
+	}
+	if set.Status.UpdatedReplicas != 0 {
+		t.Error("Failed to set updatedReplicas to 0")
+	}
 }
 
 func TestStatefulSetControl_getSetRevisions(t *testing.T) {