From beec45b4de1cd0ec639869751d124eebc276ed11 Mon Sep 17 00:00:00 2001
From: Faraaz Khan <faraaz@rationalizeit.us>
Date: Sat, 21 Apr 2018 23:15:48 -0500
Subject: [PATCH] set updated replicas

---
 .../statefulset/stateful_set_control.go       |  3 ++-
 .../statefulset/stateful_set_control_test.go  | 22 +++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/pkg/controller/statefulset/stateful_set_control.go b/pkg/controller/statefulset/stateful_set_control.go
index b1b58986043..edfc3c5bba5 100644
--- a/pkg/controller/statefulset/stateful_set_control.go
+++ b/pkg/controller/statefulset/stateful_set_control.go
@@ -297,7 +297,8 @@ func (ssc *defaultStatefulSetControl) updateStatefulSet(
 		if isCreated(pods[i]) && !isTerminating(pods[i]) {
 			if getPodRevision(pods[i]) == currentRevision.Name {
 				status.CurrentReplicas++
-			} else if getPodRevision(pods[i]) == updateRevision.Name {
+			}
+			if getPodRevision(pods[i]) == updateRevision.Name {
 				status.UpdatedReplicas++
 			}
 		}
diff --git a/pkg/controller/statefulset/stateful_set_control_test.go b/pkg/controller/statefulset/stateful_set_control_test.go
index d3a21577c26..e0de54e4288 100644
--- a/pkg/controller/statefulset/stateful_set_control_test.go
+++ b/pkg/controller/statefulset/stateful_set_control_test.go
@@ -81,6 +81,7 @@ func TestStatefulSetControl(t *testing.T) {
 		obj func() *apps.StatefulSet
 	}{
 		{CreatesPods, simpleSetFn},
+		{SetsStatusFields, simpleSetFn},
 		{ScalesUp, simpleSetFn},
 		{ScalesDown, simpleSetFn},
 		{ReplacesPods, largeSetFn},
@@ -130,6 +131,27 @@ func CreatesPods(t *testing.T, set *apps.StatefulSet, invariants invariantFunc)
 	}
 }
 
+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")
+	}
+	if set.Status.UpdatedReplicas != 3 {
+		t.Error("Failed to set UpdatedReplicas correctly")
+	}
+}
+
 func ScalesUp(t *testing.T, set *apps.StatefulSet, invariants invariantFunc) {
 	client := fake.NewSimpleClientset(set)
 	spc, _, ssc, stop := setupController(client)